| | from copy import deepcopy |
| | from typing import Dict, List, Any, Optional |
| |
|
| | import faiss |
| |
|
| | from langchain.docstore import InMemoryDocstore |
| | from langchain.embeddings import OpenAIEmbeddings |
| | from langchain.schema import Document |
| | from langchain.vectorstores import Chroma, FAISS |
| | from langchain.vectorstores.base import VectorStoreRetriever |
| |
|
| | from flows.base_flows import AtomicFlow |
| |
|
| |
|
| | class VectorStoreFlow(AtomicFlow): |
| | REQUIRED_KEYS_CONFIG = ["type", "api_keys"] |
| |
|
| | vector_db: VectorStoreRetriever |
| |
|
| | def __init__(self, vector_db, **kwargs): |
| | super().__init__(**kwargs) |
| | self.vector_db = vector_db |
| |
|
| | @classmethod |
| | def _set_up_retriever(cls, config: Dict[str, Any]) -> Dict[str, Any]: |
| | embeddings = OpenAIEmbeddings(openai_api_key=config["api_keys"]["openai"]) |
| | kwargs = {} |
| |
|
| | vs_type = config["type"] |
| |
|
| | if vs_type == "chroma": |
| | vectorstore = Chroma(config["name"], embedding_function=embeddings) |
| | elif vs_type == "faiss": |
| | index = faiss.IndexFlatL2(config.get("embedding_size", 1536)) |
| | vectorstore = FAISS( |
| | embedding_function=embeddings.embed_query, |
| | index=index, |
| | docstore=InMemoryDocstore({}), |
| | index_to_docstore_id={} |
| | ) |
| | else: |
| | raise NotImplementedError(f"Vector store '{vs_type}' not implemented") |
| |
|
| | kwargs["vector_db"] = vectorstore.as_retriever(**config.get("retriever_config", {})) |
| |
|
| | return kwargs |
| |
|
| | @classmethod |
| | def instantiate_from_config(cls, config: Dict[str, Any]): |
| | flow_config = deepcopy(config) |
| |
|
| | kwargs = {"flow_config": flow_config} |
| |
|
| | kwargs.update(cls._set_up_retriever(flow_config)) |
| |
|
| | return cls(**kwargs) |
| |
|
| | @staticmethod |
| | def package_documents(documents: List[str]) -> List[Document]: |
| | |
| | return [Document(page_content=doc, metadata={"": ""}) for doc in documents] |
| |
|
| | def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]: |
| | response = {} |
| |
|
| | operation = input_data["operation"] |
| | assert operation in ["write", "read"], f"Operation '{operation}' not supported" |
| |
|
| | content = input_data["content"] |
| | if operation == "read": |
| | assert isinstance(content, str), f"Content must be a string, got {type(content)}" |
| | query = content |
| | retrieved_documents = self.vector_db.get_relevant_documents(query) |
| | response["retrieved"] = [doc.page_content for doc in retrieved_documents] |
| | elif operation == "write": |
| | if isinstance(content, str): |
| | content = [content] |
| | assert isinstance(content, list), f"Content must be a list of strings, got {type(content)}" |
| | documents = content |
| | documents = self.package_documents(documents) |
| | self.vector_db.add_documents(documents) |
| | response["retrieved"] = "" |
| |
|
| | return response |
| |
|
