mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 23:58:34 +00:00
39 lines
1.0 KiB
Python
39 lines
1.0 KiB
Python
from typing import Any
|
|
|
|
from crewai_tools.rag.core import RAG
|
|
from crewai_tools.tools.rag.rag_tool import Adapter
|
|
|
|
|
|
class RAGAdapter(Adapter):
|
|
def __init__(
|
|
self,
|
|
collection_name: str = "crewai_knowledge_base",
|
|
persist_directory: str | None = None,
|
|
embedding_model: str = "text-embedding-3-small",
|
|
top_k: int = 5,
|
|
embedding_api_key: str | None = None,
|
|
**embedding_kwargs,
|
|
):
|
|
super().__init__()
|
|
|
|
# Prepare embedding configuration
|
|
embedding_config = {"api_key": embedding_api_key, **embedding_kwargs}
|
|
|
|
self._adapter = RAG(
|
|
collection_name=collection_name,
|
|
persist_directory=persist_directory,
|
|
embedding_model=embedding_model,
|
|
top_k=top_k,
|
|
embedding_config=embedding_config,
|
|
)
|
|
|
|
def query(self, question: str) -> str:
|
|
return self._adapter.query(question)
|
|
|
|
def add(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self._adapter.add(*args, **kwargs)
|