mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-05-02 07:42:40 +00:00
Create PDFEmbedchainAdapter class and utilize it in PDFSearchTool
This commit is contained in:
@@ -1,17 +1,23 @@
|
|||||||
from typing import Any
|
from typing import Any, Optional
|
||||||
|
|
||||||
from embedchain import App
|
from embedchain import App
|
||||||
|
|
||||||
from crewai_tools.tools.rag.rag_tool import Adapter
|
from crewai_tools.tools.rag.rag_tool import Adapter
|
||||||
|
|
||||||
|
|
||||||
class EmbedchainAdapter(Adapter):
|
class PDFEmbedchainAdapter(Adapter):
|
||||||
embedchain_app: App
|
embedchain_app: App
|
||||||
summarize: bool = False
|
summarize: bool = False
|
||||||
|
src: Optional[str] = None
|
||||||
|
|
||||||
def query(self, question: str) -> str:
|
def query(self, question: str) -> str:
|
||||||
|
where = (
|
||||||
|
{"app_id": self.embedchain_app.config.id, "source": self.src}
|
||||||
|
if self.src
|
||||||
|
else None
|
||||||
|
)
|
||||||
result, sources = self.embedchain_app.query(
|
result, sources = self.embedchain_app.query(
|
||||||
question, citations=True, dry_run=(not self.summarize)
|
question, citations=True, dry_run=(not self.summarize), where=where
|
||||||
)
|
)
|
||||||
if self.summarize:
|
if self.summarize:
|
||||||
return result
|
return result
|
||||||
@@ -22,4 +28,5 @@ class EmbedchainAdapter(Adapter):
|
|||||||
*args: Any,
|
*args: Any,
|
||||||
**kwargs: Any,
|
**kwargs: Any,
|
||||||
) -> None:
|
) -> None:
|
||||||
|
self.src = args[0] if args else None
|
||||||
self.embedchain_app.add(*args, **kwargs)
|
self.embedchain_app.add(*args, **kwargs)
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
from typing import Any, Optional, Type
|
from typing import Any, Optional, Type
|
||||||
|
|
||||||
from embedchain.models.data_type import DataType
|
from embedchain.models.data_type import DataType
|
||||||
|
from pydantic import model_validator
|
||||||
from pydantic.v1 import BaseModel, Field
|
from pydantic.v1 import BaseModel, Field
|
||||||
|
|
||||||
from ..rag.rag_tool import RagTool
|
from ..rag.rag_tool import RagTool
|
||||||
@@ -35,6 +36,22 @@ class PDFSearchTool(RagTool):
|
|||||||
self.args_schema = FixedPDFSearchToolSchema
|
self.args_schema = FixedPDFSearchToolSchema
|
||||||
self._generate_description()
|
self._generate_description()
|
||||||
|
|
||||||
|
@model_validator(mode="after")
|
||||||
|
def _set_default_adapter(self):
|
||||||
|
if isinstance(self.adapter, RagTool._AdapterPlaceholder):
|
||||||
|
from embedchain import App
|
||||||
|
|
||||||
|
from crewai_tools.adapters.pdf_embedchain_adapter import (
|
||||||
|
PDFEmbedchainAdapter,
|
||||||
|
)
|
||||||
|
|
||||||
|
app = App.from_config(config=self.config) if self.config else App()
|
||||||
|
self.adapter = PDFEmbedchainAdapter(
|
||||||
|
embedchain_app=app, summarize=self.summarize
|
||||||
|
)
|
||||||
|
|
||||||
|
return self
|
||||||
|
|
||||||
def add(
|
def add(
|
||||||
self,
|
self,
|
||||||
*args: Any,
|
*args: Any,
|
||||||
|
|||||||
Reference in New Issue
Block a user