/ src / vector_stores / protocol.py
protocol.py
 1  from __future__ import annotations
 2  
 3  """Protocol for vector store implementations."""
 4  
 5  from typing import Any, Protocol
 6  
 7  from langchain_core.documents import Document
 8  
 9  from ..constants import DEFAULT_RETRIEVAL_K
10  
11  
12  class VectorStore(Protocol):
13      """Protocol for vector store implementations.
14  
15      Any class implementing these methods is compatible with VectorStore,
16      regardless of inheritance hierarchy. This allows LangChain's vector
17      stores (Chroma, Pinecone, etc.) to work seamlessly without modification.
18  
19      """
20  
21      def similarity_search(
22          self,
23          query: str,
24          k: int = DEFAULT_RETRIEVAL_K,
25          filter: Any | None = None,
26      ) -> list[Document]:
27          """Search for similar documents.
28  
29          Parameters
30          ----------
31          query
32              Query string to search for.
33          k
34              Number of documents to retrieve.
35          filter
36              Optional metadata filter (vector store specific).
37          """
38          ...
39  
40      def similarity_search_with_score(
41          self,
42          query: str,
43          k: int = DEFAULT_RETRIEVAL_K,
44          filter: Any | None = None,
45      ) -> list[tuple[Document, float]]:
46          """Search for similar documents with similarity scores.
47  
48          Parameters
49          ----------
50          query
51              Query string to search for.
52          k
53              Number of documents to retrieve.
54          filter
55              Optional metadata filter (vector store specific).
56          """
57          ...
58  
59      def add_documents(
60          self,
61          documents: list[Document]
62      ) -> list[str]:
63          """Add documents to the vector store."""
64          ...
65  
66      def add_texts(
67          self,
68          texts: list[str],
69          metadatas: list[dict[str, Any]] | None = None,
70          ids: list[int] | None = None,
71          embeddings: list[list[float]] | None = None,
72      ) -> None:
73          """Add texts to the vector store.
74  
75          Add texts directly with optional embeddings and metadata.
76          This method is recommended for better performance when embeddings
77          are pre-computed.
78  
79          Parameters
80          ----------
81          texts
82              List of text strings to add.
83          metadatas
84              Optional list of metadata dictionaries.
85          ids
86              Optional list of document IDs.
87          embeddings
88              Optional list of pre-computed embedding vectors.
89  
90          """
91          ...
92  
93      def persist(self) -> None:
94          """Persist the vector store to disk."""
95          ...
96