mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-07-07 07:59:29 +00:00
66 lines
2.0 KiB
Python
66 lines
2.0 KiB
Python
from collections.abc import Callable
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from crewai.utilities.lock_store import lock as store_lock
|
|
from lancedb import ( # type: ignore[import-untyped]
|
|
connect as lancedb_connect,
|
|
)
|
|
from openai import Client as OpenAIClient
|
|
from pydantic import Field, PrivateAttr
|
|
|
|
from crewai_tools.tools.rag.rag_tool import Adapter
|
|
|
|
|
|
def _default_embedding_function() -> Callable[[list[str]], list[list[float]]]:
|
|
"""Create a default embedding function using OpenAI."""
|
|
client = OpenAIClient()
|
|
|
|
def _embedding_function(input: list[str]) -> list[list[float]]:
|
|
rs = client.embeddings.create(input=input, model="text-embedding-ada-002")
|
|
return [record.embedding for record in rs.data]
|
|
|
|
return _embedding_function
|
|
|
|
|
|
class LanceDBAdapter(Adapter):
|
|
uri: str | Path
|
|
table_name: str
|
|
embedding_function: Callable[[list[str]], list[list[float]]] = Field(
|
|
default_factory=_default_embedding_function
|
|
)
|
|
top_k: int = 3
|
|
vector_column_name: str = "vector"
|
|
text_column_name: str = "text"
|
|
|
|
_db: Any = PrivateAttr()
|
|
_table: Any = PrivateAttr()
|
|
_lock_name: str = PrivateAttr(default="")
|
|
|
|
def model_post_init(self, __context: Any) -> None:
|
|
self._db = lancedb_connect(self.uri)
|
|
self._table = self._db.open_table(self.table_name)
|
|
self._lock_name = f"lancedb:{os.path.realpath(str(self.uri))}"
|
|
|
|
super().model_post_init(__context)
|
|
|
|
def query(self, question: str) -> str: # type: ignore[override]
|
|
query = self.embedding_function([question])[0]
|
|
results = (
|
|
self._table.search(query, vector_column_name=self.vector_column_name)
|
|
.limit(self.top_k)
|
|
.select([self.text_column_name])
|
|
.to_list()
|
|
)
|
|
values = [result[self.text_column_name] for result in results]
|
|
return "\n".join(values)
|
|
|
|
def add(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
with store_lock(self._lock_name):
|
|
self._table.add(*args, **kwargs)
|