diff --git a/src/poetry.lock b/src/poetry.lock index e1bd8887..a408225a 100644 --- a/src/poetry.lock +++ b/src/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.4.1 and should not be changed by hand. [[package]] name = "aiohappyeyeballs" @@ -1624,8 +1624,8 @@ google-api-core = {version = ">=2.17.1,<3.0.0", extras = ["grpc"]} google-auth = ">=2.14.1,<2.24.0 || >2.24.0,<2.25.0 || >2.25.0,<3.0.0" grpcio = ">=1.59.0,<2.0.0" proto-plus = [ - {version = ">=1.25.0,<2.0.0", markers = "python_version >= \"3.13\""}, {version = ">=1.22.3,<2.0.0", markers = "python_version < \"3.13\""}, + {version = ">=1.25.0,<2.0.0", markers = "python_version >= \"3.13\""}, ] protobuf = ">=4.25.8,<8.0.0" @@ -2414,7 +2414,7 @@ files = [ [package.dependencies] attrs = ">=22.2.0" -jsonschema-specifications = ">=2023.03.6" +jsonschema-specifications = ">=2023.3.6" referencing = ">=0.28.4" rpds-py = ">=0.25.0" @@ -2451,7 +2451,7 @@ files = [ [package.dependencies] aiohttp = ">=3.13.5,<4.0.0" -certifi = ">=14.05.14" +certifi = ">=14.5.14" durationpy = ">=0.7" python-dateutil = ">=2.5.3" pyyaml = ">=6.0.3" @@ -2481,26 +2481,6 @@ anthropic = ">=0.120.0,<1.0.0" langchain-core = ">=1.5.2,<2.0.0" pydantic = ">=2.7.4,<3.0.0" -[[package]] -name = "langchain-chroma" -version = "0.2.6" -description = "An integration package connecting Chroma and LangChain." -optional = false -python-versions = ">=3.9" -groups = ["optional"] -files = [ - {file = "langchain_chroma-0.2.6-py3-none-any.whl", hash = "sha256:d7e10101b0942cd990eedb798c3d85ed3e8415a992c8a388843196f6ab97b41b"}, - {file = "langchain_chroma-0.2.6.tar.gz", hash = "sha256:ec5ca0f6f7692ac053741e076ea086c4be0cfcb5846c8693b1bcc3089c88b65e"}, -] - -[package.dependencies] -chromadb = ">=1.0.20" -langchain-core = ">=0.3.76" -numpy = [ - {version = ">=2.1.0", markers = "python_version >= \"3.13\""}, - {version = ">=1.26.0", markers = "python_version < \"3.13\""}, -] - [[package]] name = "langchain-core" version = "1.5.3" @@ -2797,7 +2777,7 @@ colorama = {version = ">=0.3.4", markers = "sys_platform == \"win32\""} win32-setctime = {version = ">=1.0.0", markers = "sys_platform == \"win32\""} [package.extras] -dev = ["Sphinx (==8.1.3) ; python_version >= \"3.11\"", "build (==1.2.2) ; python_version >= \"3.11\"", "colorama (==0.4.5) ; python_version < \"3.8\"", "colorama (==0.4.6) ; python_version >= \"3.8\"", "exceptiongroup (==1.1.3) ; python_version >= \"3.7\" and python_version < \"3.11\"", "freezegun (==1.1.0) ; python_version < \"3.8\"", "freezegun (==1.5.0) ; python_version >= \"3.8\"", "mypy (==v0.910) ; python_version < \"3.6\"", "mypy (==v0.971) ; python_version == \"3.6\"", "mypy (==v1.13.0) ; python_version >= \"3.8\"", "mypy (==v1.4.1) ; python_version == \"3.7\"", "myst-parser (==4.0.0) ; python_version >= \"3.11\"", "pre-commit (==4.0.1) ; python_version >= \"3.9\"", "pytest (==6.1.2) ; python_version < \"3.8\"", "pytest (==8.3.2) ; python_version >= \"3.8\"", "pytest-cov (==2.12.1) ; python_version < \"3.8\"", "pytest-cov (==5.0.0) ; python_version == \"3.8\"", "pytest-cov (==6.0.0) ; python_version >= \"3.9\"", "pytest-mypy-plugins (==1.9.3) ; python_version >= \"3.6\" and python_version < \"3.8\"", "pytest-mypy-plugins (==3.1.0) ; python_version >= \"3.8\"", "sphinx-rtd-theme (==3.0.2) ; python_version >= \"3.11\"", "tox (==3.27.1) ; python_version < \"3.8\"", "tox (==4.23.2) ; python_version >= \"3.8\"", "twine (==6.0.1) ; python_version >= \"3.11\""] +dev = ["Sphinx (==8.1.3) ; python_version >= \"3.11\"", "build (==1.2.2) ; python_version >= \"3.11\"", "colorama (==0.4.5) ; python_version < \"3.8\"", "colorama (==0.4.6) ; python_version >= \"3.8\"", "exceptiongroup (==1.1.3) ; python_version >= \"3.7\" and python_version < \"3.11\"", "freezegun (==1.1.0) ; python_version < \"3.8\"", "freezegun (==1.5.0) ; python_version >= \"3.8\"", "mypy (==0.910) ; python_version < \"3.6\"", "mypy (==0.971) ; python_version == \"3.6\"", "mypy (==1.13.0) ; python_version >= \"3.8\"", "mypy (==1.4.1) ; python_version == \"3.7\"", "myst-parser (==4.0.0) ; python_version >= \"3.11\"", "pre-commit (==4.0.1) ; python_version >= \"3.9\"", "pytest (==6.1.2) ; python_version < \"3.8\"", "pytest (==8.3.2) ; python_version >= \"3.8\"", "pytest-cov (==2.12.1) ; python_version < \"3.8\"", "pytest-cov (==5.0.0) ; python_version == \"3.8\"", "pytest-cov (==6.0.0) ; python_version >= \"3.9\"", "pytest-mypy-plugins (==1.9.3) ; python_version >= \"3.6\" and python_version < \"3.8\"", "pytest-mypy-plugins (==3.1.0) ; python_version >= \"3.8\"", "sphinx-rtd-theme (==3.0.2) ; python_version >= \"3.11\"", "tox (==3.27.1) ; python_version < \"3.8\"", "tox (==4.23.2) ; python_version >= \"3.8\"", "twine (==6.0.1) ; python_version >= \"3.11\""] [[package]] name = "markdown-it-py" @@ -3803,8 +3783,8 @@ files = [ [package.dependencies] googleapis-common-protos = ">=1.57,<2.0" grpcio = [ - {version = ">=1.66.2,<2.0.0", markers = "python_version == \"3.13\""}, {version = ">=1.63.2,<2.0.0", markers = "python_version < \"3.13\""}, + {version = ">=1.66.2,<2.0.0", markers = "python_version == \"3.13\""}, ] opentelemetry-api = ">=1.15,<2.0" opentelemetry-exporter-otlp-proto-common = "1.44.0" @@ -5599,10 +5579,10 @@ files = [ ] [package.dependencies] -botocore = ">=1.37.4,<2.0a.0" +botocore = ">=1.37.4,<2.0a0" [package.extras] -crt = ["botocore[crt] (>=1.37.4,<2.0a.0)"] +crt = ["botocore[crt] (>=1.37.4,<2.0a0)"] [[package]] name = "safetensors" @@ -7746,9 +7726,9 @@ files = [ ] [package.extras] -cffi = ["cffi (>=1.17,<2.0) ; platform_python_implementation != \"PyPy\" and python_version < \"3.14\"", "cffi (>=2.0.0b) ; platform_python_implementation != \"PyPy\" and python_version >= \"3.14\""] +cffi = ["cffi (>=1.17,<2.0) ; platform_python_implementation != \"PyPy\" and python_version < \"3.14\"", "cffi (>=2.0.0b0) ; platform_python_implementation != \"PyPy\" and python_version >= \"3.14\""] [metadata] lock-version = "2.1" python-versions = "<3.14,>=3.10" -content-hash = "8b842dbeea201850311085956528916b6eb950e0638354c15f9f4242a3a11b28" +content-hash = "808a9d00aa11a751f27a9a8ee552bc1098966814785d3f425b0c12ef2b71ba9c" diff --git a/src/pyproject.toml b/src/pyproject.toml index 1e2c69bb..47db04f9 100644 --- a/src/pyproject.toml +++ b/src/pyproject.toml @@ -79,7 +79,6 @@ en_core_web_sm = {url = "https://github.com/explosion/spacy-models/releases/down # requires running the Chroma server with trust_remote_code=true, which Sherpa # does not do. Bump once an upstream fixed release ships. chromadb = "^1.0.9" -langchain-chroma = "^0.2.5" langchain-anthropic = ">=1.0.0,<2.0" langchain-google-genai = ">=2.0.0,<3.0" boto3 = "^1.28.77" diff --git a/src/sherpa_ai/connectors/scripts/query_chroma.py b/src/sherpa_ai/connectors/scripts/query_chroma.py index 5b34f5d5..0b3b551e 100644 --- a/src/sherpa_ai/connectors/scripts/query_chroma.py +++ b/src/sherpa_ai/connectors/scripts/query_chroma.py @@ -2,10 +2,10 @@ import json import uuid -import chromadb -from chromadb.config import Settings -from dotenv import load_dotenv -from langchain_openai import OpenAIEmbeddings +import chromadb +from chromadb.config import Settings +from chromadb.utils import embedding_functions +from dotenv import load_dotenv from loguru import logger @@ -16,29 +16,17 @@ def main(args): settings=Settings(allow_reset=True), ) - embedding_func = OpenAIEmbeddings() - try: - from langchain_chroma import Chroma - except ImportError: - raise ImportError( - "Could not import langchain_chroma python package. " - "This is needed in order to use Chroma. " - "Please install it with `pip install langchain-chroma`" - ) - chroma = Chroma( - client=client, - collection_name=args.chroma_index, - embedding_function=embedding_func, + embedding_func = embedding_functions.OpenAIEmbeddingFunction( + model_name="text-embedding-ada-002" + ) + collection = client.get_or_create_collection( + name=args.chroma_index, embedding_function=embedding_func ) query = input("Enter query: ") - results = chroma.similarity_search( - query=query, - number_of_results=5, - k=1 - ) + results = collection.query(query_texts=[query], n_results=1) - logger.info(results[0].page_content) + logger.info(results["documents"][0][0]) logger.info("Done! Chroma is up and running.") diff --git a/src/sherpa_ai/connectors/vectorstores.py b/src/sherpa_ai/connectors/vectorstores.py index 40d1d9fe..7f91c814 100644 --- a/src/sherpa_ai/connectors/vectorstores.py +++ b/src/sherpa_ai/connectors/vectorstores.py @@ -1,5 +1,10 @@ import os +import uuid +from typing import Any, Iterable, List, Optional, Tuple, Type +from langchain_core.documents import Document +from langchain_core.embeddings import Embeddings +from langchain_core.vectorstores import VectorStore, VectorStoreRetriever from langchain_openai import OpenAIEmbeddings from langchain_text_splitters import CharacterTextSplitter from loguru import logger @@ -8,32 +13,413 @@ from sherpa_ai.utils import load_files -class LocalChromaStore: - """A local Chroma-based vector store. +class ConversationStore(VectorStore): + """A vector store for storing and retrieving conversation data. - This class extends the Chroma vector store to provide additional functionality - for working with local files. + This class provides methods to store conversation data in a vector database + and retrieve similar conversations based on queries. + + Attributes: + db: The underlying database connection. + namespace (str): The namespace for the vector store. + embeddings_func: The embedding function to use. + text_key (str): The key used to store the text in metadata. + + Example: + >>> from sherpa_ai.connectors.vectorstores import ConversationStore + >>> store = ConversationStore.from_index("my_namespace", "api_key", "my_index") + >>> store.add_text("This is a conversation", {"user": "user1"}) + >>> results = store.similarity_search("conversation", top_k=5) + """ + def __init__(self, namespace, db, embeddings, text_key): + """Initialize a ConversationStore instance. + + Args: + namespace (str): The namespace for the vector store. + db: The database connection. + embeddings: The embedding function to use. + text_key (str): The key used to store the text in metadata. + + Example: + >>> from sherpa_ai.connectors.vectorstores import ConversationStore + >>> store = ConversationStore("my_namespace", db, embeddings, "text") + """ + self.db = db + self.namespace = namespace + self.embeddings_func = embeddings + self.text_key = text_key + + @classmethod + def from_index(cls, namespace, openai_api_key, index_name, text_key="text"): + """Create a ConversationStore from a Pinecone index. + + This method initializes a Pinecone client and creates a ConversationStore + instance connected to the specified index. + + Args: + namespace (str): The namespace for the vector store. + openai_api_key (str): The OpenAI API key. + index_name (str): The name of the Pinecone index. + text_key (str, optional): The key used to store the text in metadata. Defaults to "text". + + Returns: + ConversationStore: A new ConversationStore instance. + + Raises: + ImportError: If the pinecone-client package is not installed. + + Example: + >>> from sherpa_ai.connectors.vectorstores import ConversationStore + >>> store = ConversationStore.from_index("my_namespace", "api_key", "my_index") + """ + try: + import pinecone + except ImportError: + raise ImportError( + "Could not import pinecone-client python package. " + "This is needed in order to to use ConversationStore. " + "Please install it with `pip install pinecone-client`" + ) + + pinecone.init(api_key=cfg.PINECONE_API_KEY, environment=cfg.PINECONE_ENV) + logger.info(f"Loading index {index_name} from Pinecone") + index = pinecone.Index(index_name) + embedding = OpenAIEmbeddings(openai_api_key=openai_api_key) + return cls(namespace, index, embedding, text_key) + + def add_text(self, text: str, metadata={}) -> str: + """Add a single text to the vector store. + + This method embeds the text, adds it to the database with the provided metadata, + and returns the ID of the added text. + + Args: + text (str): The text to add. + metadata (dict, optional): Metadata to associate with the text. Defaults to {}. + + Returns: + str: The ID of the added text. + + Example: + >>> from sherpa_ai.connectors.vectorstores import ConversationStore + >>> store = ConversationStore.from_index("my_namespace", "api_key", "my_index") + >>> id = store.add_text("This is a conversation", {"user": "user1"}) + >>> print(id) + '123e4567-e89b-12d3-a456-426614174000' + """ + metadata[self.text_key] = text + id = str(uuid.uuid4()) + embedding = self.embeddings.embed_query(text) + doc = {"id": id, "values": embedding, "metadata": metadata} + self.db.upsert(vectors=[doc], namespace=self.namespace) + + return id + + @property + def embeddings(self) -> Optional[Embeddings]: + """Access the query embedding object if available.""" + return self.embeddings_func + + def add_texts(self, texts: Iterable[str], metadatas: List[dict]) -> List[str]: + """Add multiple texts to the vector store. + + This method adds each text with its corresponding metadata to the vector store. + + Args: + texts (Iterable[str]): The texts to add. + metadatas (List[dict]): The metadata for each text. + + Example: + >>> from sherpa_ai.connectors.vectorstores import ConversationStore + >>> store = ConversationStore.from_index("my_namespace", "api_key", "my_index") + >>> texts = ["Text 1", "Text 2"] + >>> metadatas = [{"user": "user1"}, {"user": "user2"}] + >>> store.add_texts(texts, metadatas) + """ + for text, metadata in zip(texts, metadatas): + self.add_text(text, metadata) + + def similarity_search( + self, + text: str, + top_k: int = 5, + filter: Optional[dict] = None, + threshold: float = 0.7, + ) -> list[Document]: + """Perform a similarity search in the vector store. + + This method searches for texts that are semantically similar to the query. + + Args: + text (str): The search query. + top_k (int, optional): The number of results to return. Defaults to 5. + filter (Optional[dict], optional): Filter criteria for the search. Defaults to None. + threshold (float, optional): The similarity threshold. Defaults to 0.7. + + Returns: + list[Document]: A list of documents that match the query. + + Example: + >>> from sherpa_ai.connectors.vectorstores import ConversationStore + >>> store = ConversationStore.from_index("my_namespace", "api_key", "my_index") + >>> results = store.similarity_search("What is machine learning?", top_k=5) + >>> for doc in results: + ... print(doc.page_content[:100]) + """ + query_embedding = self.embeddings.embed_query(text) + results = self.db.query( + [query_embedding], + top_k=top_k, + include_metadata=True, + namespace=self.namespace, + filter=filter, + ) + + docs = [] + for res in results["matches"]: + metadata = res["metadata"] + text = metadata.pop(self.text_key) + if res["score"] > threshold: + docs.append(Document(page_content=text, metadata=metadata)) + return docs + + def _similarity_search_with_relevance_scores( + self, + query: str, + k: int = 4, + **kwargs: Any, + ) -> List[Tuple[Document, float]]: + """Perform a similarity search and return documents with relevance scores. + + This method searches for texts that are semantically similar to the query + and returns them along with their relevance scores. + + Args: + query (str): The search query. + k (int, optional): The number of results to return. Defaults to 4. + **kwargs: Additional keyword arguments. + + Returns: + List[Tuple[Document, float]]: A list of tuples containing documents and their relevance scores. + + Example: + >>> from sherpa_ai.connectors.vectorstores import ConversationStore + >>> store = ConversationStore.from_index("my_namespace", "api_key", "my_index") + >>> results = store._similarity_search_with_relevance_scores("What is machine learning?") + >>> for doc, score in results: + ... print(f"Score: {score}, Content: {doc.page_content[:100]}") + """ + logger.debug("query", query) + query_embedding = self.embeddings.embed_query(query) + results = self.db.query( + [query_embedding], + top_k=k, + include_metadata=True, + namespace=self.namespace, + filter=kwargs.get("filter", None), + ) + + docs_with_score = [] + for res in results["matches"]: + metadata = res["metadata"] + text = metadata.pop(self.text_key) + docs_with_score.append( + (Document(page_content=text, metadata=metadata), res["score"]) + ) + logger.debug(docs_with_score) + return docs_with_score + + @classmethod + def delete(cls, namespace, index_name): + """Delete all vectors in a namespace. + + This method deletes all vectors in the specified namespace of the Pinecone index. + + Args: + namespace (str): The namespace to delete. + index_name (str): The name of the Pinecone index. + + Returns: + The result of the delete operation. + + Raises: + ImportError: If the pinecone-client package is not installed. + + Example: + >>> from sherpa_ai.connectors.vectorstores import ConversationStore + >>> ConversationStore.delete("my_namespace", "my_index") + """ + try: + import pinecone + except ImportError: + raise ImportError( + "Could not import pinecone-client python package. " + "This is needed in order to to use ConversationStore. " + "Please install it with `pip install pinecone-client`" + ) + + + pinecone.init(api_key=cfg.PINECONE_API_KEY, environment=cfg.PINECONE_ENV) + index = pinecone.Index(index_name) + return index.delete(delete_all=True, namespace=namespace) + + @classmethod + def get_vector_retrieval( + cls, + namespace: str, + openai_api_key: str, + index_name: str, + search_type="similarity", + search_kwargs={}, + ) -> VectorStoreRetriever: + """Create a vector store retriever. + + This method creates a ConversationStore and returns a VectorStoreRetriever + for it. + + Args: + namespace (str): The namespace for the vector store. + openai_api_key (str): The OpenAI API key. + index_name (str): The name of the Pinecone index. + search_type (str, optional): The type of search to perform. Defaults to "similarity". + search_kwargs (dict, optional): Additional keyword arguments for the search. Defaults to {}. + + Returns: + VectorStoreRetriever: A retriever for the vector store. + + Example: + >>> from sherpa_ai.connectors.vectorstores import ConversationStore + >>> retriever = ConversationStore.get_vector_retrieval("my_namespace", "api_key", "my_index") + >>> results = retriever.get_relevant_documents("What is machine learning?") + """ + vectorstore = cls.from_index(namespace, openai_api_key, index_name) + retriever = VectorStoreRetriever( + vectorstore=vectorstore, + search_type=search_type, + search_kwargs=search_kwargs, + ) + return retriever + + @classmethod + def from_texts( + cls, + texts: List[str], + embedding: Embeddings, + metadatas: Optional[list[dict]] = None, + namespace: str = "", + db: Optional[Any] = None, + text_key: str = "text", + **kwargs: Any, + ) -> "ConversationStore": + """Create a ConversationStore from a list of texts. + + Args: + texts (List[str]): The texts to add. + embedding (Embeddings): The embedding function to use. + metadatas (Optional[list[dict]], optional): The metadata for each text. + Defaults to an empty dict per text. + namespace (str, optional): The namespace for the vector store. + db: The database connection (e.g. a Pinecone index) to store vectors in. + text_key (str, optional): The key used to store the text in metadata. + Defaults to "text". + + Returns: + ConversationStore: A new ConversationStore instance with the texts added. + + Example: + >>> from sherpa_ai.connectors.vectorstores import ConversationStore + >>> store = ConversationStore.from_texts( + ... ["hello", "world"], embedding, namespace="my_namespace", db=index + ... ) + """ + store = cls(namespace, db, embedding, text_key) + metadatas = metadatas or [{} for _ in texts] + store.add_texts(texts, metadatas) + return store + + +class LocalChromaStore(VectorStore): + """A local Chroma-based vector store, backed directly by chromadb. + + chromadb is already a hard dependency of sherpa-ai (see chroma_vector_store.py), + so this talks to it directly instead of going through the langchain-chroma + wrapper package, which brings nothing this class needs. Example: >>> from sherpa_ai.connectors.vectorstores import LocalChromaStore >>> store = LocalChromaStore.from_folder("path/to/files", "api_key") >>> results = store.similarity_search("query", k=5) """ - - def __init__(self, *args, **kwargs): + + def __init__( + self, + collection_name: str = "langchain", + embedding_function: Optional[Embeddings] = None, + client: Optional[Any] = None, + ): try: - from langchain_chroma import Chroma + import chromadb except ImportError: raise ImportError( - "Could not import langchain_chroma python package. " + "Could not import chromadb python package. " "This is needed in order to use LocalChromaStore. " - "Please install it with `pip install langchain-chroma`" + "Please install it with `pip install chromadb`" ) - self._chroma = Chroma(*args, **kwargs) - - def __getattr__(self, name): - """Delegate attribute access to the underlying Chroma instance.""" - return getattr(self._chroma, name) + + self._embedding_function = embedding_function + self._client = client if client is not None else chromadb.EphemeralClient() + self._collection = self._client.get_or_create_collection(name=collection_name) + + @property + def embeddings(self) -> Optional[Embeddings]: + """Access the query embedding object if available.""" + return self._embedding_function + + def add_texts( + self, + texts: Iterable[str], + metadatas: Optional[List[dict]] = None, + **kwargs: Any, + ) -> List[str]: + """Embed and add texts to the underlying chromadb collection.""" + texts = list(texts) + ids = [str(uuid.uuid4()) for _ in texts] + embeddings = self._embedding_function.embed_documents(texts) + self._collection.add( + ids=ids, + embeddings=embeddings, + documents=texts, + metadatas=metadatas if metadatas else [{} for _ in texts], + ) + return ids + + def similarity_search(self, query: str, k: int = 4, **kwargs: Any) -> List[Document]: + """Perform a similarity search in the chromadb collection.""" + query_embedding = self._embedding_function.embed_query(query) + results = self._collection.query(query_embeddings=[query_embedding], n_results=k) + + docs = [] + documents = results.get("documents") or [[]] + metadatas = results.get("metadatas") or [[]] + for text, metadata in zip(documents[0], metadatas[0]): + docs.append(Document(page_content=text, metadata=metadata or {})) + return docs + + @classmethod + def from_texts( + cls, + texts: List[str], + embedding: Embeddings, + metadatas: Optional[List[dict]] = None, + index_name: str = "langchain", + **kwargs: Any, + ) -> "LocalChromaStore": + """Create a LocalChromaStore from a list of texts.""" + store = cls(collection_name=index_name, embedding_function=embedding) + if texts: + store.add_texts(texts, metadatas) + return store + @classmethod def from_folder(cls, file_path, openai_api_key, index_name="chroma"): """Create a Chroma DB from a folder of files. @@ -97,37 +483,19 @@ def configure_chroma(host: str, port: int, index_name: str, openai_api_key: str) "This is needed in order to to use Chroma. " "Please install it with `pip install chromadb" ) - - try: - from langchain_chroma import Chroma - except ImportError: - raise ImportError( - "Could not import langchain_chroma python package. " - "This is needed in order to use Chroma. " - "Please install it with `pip install langchain-chroma`" - ) + client = chromadb.HttpClient(host=cfg.CHROMA_HOST, port=cfg.CHROMA_PORT) embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key) - chroma = Chroma( - client=client, collection_name=cfg.CHROMA_INDEX, embedding_function=embeddings + return LocalChromaStore( + collection_name=cfg.CHROMA_INDEX, embedding_function=embeddings, client=client ) - return chroma - - -def _is_chroma_available(): - """Check if langchain_chroma is available.""" - try: - import langchain_chroma - return True - except ImportError: - return False def get_vectordb(): """Get a vector database retriever based on configuration. This function returns a vector database retriever based on the configuration - in the config module. It supports Chroma and local ChromaDB. + in the config module. It supports Pinecone, Chroma, and local ChromaDB. Returns: VectorStoreRetriever: A retriever for the vector store. @@ -137,20 +505,19 @@ def get_vectordb(): >>> retriever = get_vectordb() >>> results = retriever.get_relevant_documents("What is machine learning?") """ - if cfg.VECTORDB == "chroma": + if cfg.VECTORDB == "pinecone": + return ConversationStore.get_vector_retrieval( + cfg.PINECONE_NAMESPACE, + cfg.OPENAI_API_KEY, + index_name=cfg.PINECONE_INDEX, + search_type="similarity_score_threshold", + search_kwargs={"score_threshold": 0.0}, + ) + elif cfg.VECTORDB == "chroma": return configure_chroma( cfg.CHROMA_HOST, cfg.CHROMA_PORT, cfg.CHROMA_INDEX, cfg.OPENAI_API_KEY ).as_retriever() else: - # Check if langchain_chroma is available before trying to use it - if not _is_chroma_available(): - raise ImportError( - "Could not import langchain_chroma python package. " - "This is needed in order to use the default vector store. " - "Please install it with `pip install langchain-chroma` or " - "configure a different vector database (chroma) in your environment." - ) - if os.path.exists("files"): return LocalChromaStore.from_folder( "files", cfg.OPENAI_API_KEY diff --git a/src/tests/unit_tests/actions/test_context_search.py b/src/tests/unit_tests/actions/test_context_search.py index 24c3100b..6ff8474e 100644 --- a/src/tests/unit_tests/actions/test_context_search.py +++ b/src/tests/unit_tests/actions/test_context_search.py @@ -6,15 +6,15 @@ def _is_chroma_available(): - """Check if langchain_chroma is available.""" + """Check if chromadb is available.""" try: - import langchain_chroma + import chromadb return True except ImportError: return False -# Only import ContextSearch if langchain_chroma is available +# Only import ContextSearch if chromadb is available if _is_chroma_available(): from sherpa_ai.actions.context_search import ContextSearch @@ -36,7 +36,7 @@ def mock_context_search(external_api): yield -@pytest.mark.skipif(not _is_chroma_available(), reason="langchain_chroma not available") +@pytest.mark.skipif(not _is_chroma_available(), reason="chromadb not available") def test_context_search_succeeds(get_llm, mock_context_search): # noqa: F811 role_description = ( "The programmer receives requirements about a program and write it" diff --git a/src/tests/unit_tests/connectors/test_conversation_store.py b/src/tests/unit_tests/connectors/test_conversation_store.py new file mode 100644 index 00000000..8a7e98bf --- /dev/null +++ b/src/tests/unit_tests/connectors/test_conversation_store.py @@ -0,0 +1,20 @@ +from unittest.mock import MagicMock + +from sherpa_ai.connectors.vectorstores import ConversationStore + + +def test_from_texts_creates_store_and_adds_texts(): + fake_db = MagicMock() + embedding = MagicMock() + embedding.embed_query.side_effect = [[0.1], [0.2]] + + store = ConversationStore.from_texts( + ["hello", "world"], + embedding, + metadatas=[{"a": 1}, {"b": 2}], + namespace="ns", + db=fake_db, + ) + + assert isinstance(store, ConversationStore) + assert fake_db.upsert.call_count == 2 diff --git a/src/tests/unit_tests/connectors/test_local_chroma_store.py b/src/tests/unit_tests/connectors/test_local_chroma_store.py new file mode 100644 index 00000000..0939204b --- /dev/null +++ b/src/tests/unit_tests/connectors/test_local_chroma_store.py @@ -0,0 +1,62 @@ +import builtins + +import pytest + +from sherpa_ai.connectors.vectorstores import LocalChromaStore + + +class FakeEmbeddings: + """Deterministic stand-in for OpenAIEmbeddings: each text maps to a + fixed-size vector based on its length, so unrelated texts don't collide.""" + + def embed_documents(self, texts): + return [self.embed_query(t) for t in texts] + + def embed_query(self, text): + return [float(len(text)), float(sum(ord(c) for c in text) % 97)] + + +@pytest.fixture +def store(): + return LocalChromaStore(collection_name="test", embedding_function=FakeEmbeddings()) + + +def test_init_raises_clear_error_when_chromadb_missing(monkeypatch): + real_import = builtins.__import__ + + def fake_import(name, *args, **kwargs): + if name == "chromadb": + raise ImportError("no chromadb") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + with pytest.raises(ImportError, match="Could not import chromadb"): + LocalChromaStore(collection_name="test", embedding_function=FakeEmbeddings()) + + +def test_add_texts_and_similarity_search_round_trip(store): + store.add_texts( + ["sherpa helps you climb mountains", "bananas are yellow"], + metadatas=[{"topic": "sherpa"}, {"topic": "fruit"}], + ) + + results = store.similarity_search("sherpa helps you climb mountains", k=1) + + assert len(results) == 1 + assert results[0].page_content == "sherpa helps you climb mountains" + assert results[0].metadata["topic"] == "sherpa" + + +def test_from_texts_builds_a_queryable_store(): + store = LocalChromaStore.from_texts( + ["sherpa helps you climb mountains", "bananas are yellow"], + embedding=FakeEmbeddings(), + metadatas=[{"topic": "sherpa"}, {"topic": "fruit"}], + index_name="test_from_texts", + ) + + results = store.similarity_search("bananas are yellow", k=1) + + assert len(results) == 1 + assert results[0].page_content == "bananas are yellow"