mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 00:28:31 +00:00
33 lines
865 B
Python
33 lines
865 B
Python
from typing import Any, Optional
|
|
|
|
from embedchain import App
|
|
|
|
from crewai_tools.tools.rag.rag_tool import Adapter
|
|
|
|
|
|
class PDFEmbedchainAdapter(Adapter):
|
|
embedchain_app: App
|
|
summarize: bool = False
|
|
src: Optional[str] = None
|
|
|
|
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(
|
|
question, citations=True, dry_run=(not self.summarize), where=where
|
|
)
|
|
if self.summarize:
|
|
return result
|
|
return "\n\n".join([source[0] for source in sources])
|
|
|
|
def add(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self.src = args[0] if args else None
|
|
self.embedchain_app.add(*args, **kwargs)
|