| """Retriever that generates and executes structured queries over its own data source. |
| |
| NOTE: This code is adapted from the original implementation in the LangChain repo, |
| but has been modified to work with the KTH QA system. |
| |
| """ |
|
|
| from langchain.vectorstores import Pinecone, VectorStore |
| from langchain.schema import BaseRetriever, Document |
| from langchain.retrievers.self_query.pinecone import PineconeTranslator |
| from langchain.chains.query_constructor.schema import AttributeInfo |
| from langchain.chains.query_constructor.ir import StructuredQuery, Visitor |
| from langchain.chains.query_constructor.base import load_query_constructor_chain |
| from langchain.base_language import BaseLanguageModel |
| from langchain import LLMChain |
| from pydantic import BaseModel, Field, root_validator |
| import re |
| from typing import Any, Dict, List, Optional, Type, cast |
| import logging |
| logger = logging.getLogger() |
|
|
|
|
| COURSE_PATTERN = r"[a-zA-Z]{2,3}\d{3,4}\w?" |
|
|
|
|
| def make_uppercase(matchobj): |
| return matchobj.group(0).upper() |
|
|
|
|
| def _get_builtin_translator(vectorstore_cls: Type[VectorStore]) -> Visitor: |
| """Get the translator class corresponding to the vector store class.""" |
| BUILTIN_TRANSLATORS: Dict[Type[VectorStore], Type[Visitor]] = { |
| Pinecone: PineconeTranslator |
| } |
| if vectorstore_cls not in BUILTIN_TRANSLATORS: |
| raise ValueError( |
| f"Self query retriever with Vector Store type {vectorstore_cls}" |
| f" not supported." |
| ) |
| return BUILTIN_TRANSLATORS[vectorstore_cls]() |
|
|
|
|
| class SelfQueryRetriever(BaseRetriever, BaseModel): |
| """Retriever that wraps around a vector store and uses an LLM to generate |
| the vector store queries.""" |
|
|
| vectorstore: VectorStore |
| """The underlying vector store from which documents will be retrieved.""" |
| llm_chain: LLMChain |
| """The LLMChain for generating the vector store queries.""" |
| search_type: str = "similarity" |
| """The search type to perform on the vector store.""" |
| search_kwargs: dict = Field(default_factory=dict) |
| """Keyword arguments to pass in to the vector store search.""" |
| structured_query_translator: Visitor |
| """Translator for turning internal query language into vectorstore search params.""" |
| verbose: bool = False |
|
|
| class Config: |
| """Configuration for this pydantic object.""" |
|
|
| arbitrary_types_allowed = True |
|
|
| @root_validator(pre=True) |
| def validate_translator(cls, values: Dict) -> Dict: |
| """Validate translator.""" |
| if "structured_query_translator" not in values: |
| vectorstore_cls = values["vectorstore"].__class__ |
| values["structured_query_translator"] = _get_builtin_translator( |
| vectorstore_cls |
| ) |
| return values |
|
|
| def get_relevant_documents(self, query: str) -> List[Document]: |
| """Get documents relevant for a query. |
| |
| Args: |
| query: string to find relevant documents for |
| |
| Returns: |
| List of relevant documents |
| """ |
| if re.findall(COURSE_PATTERN, query): |
| query = re.sub(COURSE_PATTERN, make_uppercase, query) |
| inputs = self.llm_chain.prep_inputs(query) |
| structured_query = cast( |
| StructuredQuery, self.llm_chain.predict_and_parse( |
| callbacks=None, **inputs) |
| ) |
| if self.verbose: |
| logger.info( |
| "Found course pattern in query, using structured query:") |
| logger.info(structured_query) |
| new_query, new_kwargs = self.structured_query_translator.visit_structured_query( |
| structured_query |
| ) |
| search_kwargs = {**self.search_kwargs, **new_kwargs} |
| else: |
| search_kwargs = self.search_kwargs |
| docs = self.vectorstore.search( |
| query, self.search_type, **search_kwargs) |
| return docs |
| |
| async def aget_relevant_documents(self, query: str) -> List[Document]: |
| raise NotImplementedError |
|
|
| @classmethod |
| def from_llm( |
| cls, |
| llm: BaseLanguageModel, |
| vectorstore: VectorStore, |
| document_contents: str, |
| metadata_field_info: List[AttributeInfo], |
| structured_query_translator: Optional[Visitor] = None, |
| chain_kwargs: Optional[Dict] = None, |
| **kwargs: Any, |
| ) -> "SelfQueryRetriever": |
| if structured_query_translator is None: |
| structured_query_translator = _get_builtin_translator( |
| vectorstore.__class__) |
| chain_kwargs = chain_kwargs or {} |
| if "allowed_comparators" not in chain_kwargs: |
| chain_kwargs[ |
| "allowed_comparators" |
| ] = structured_query_translator.allowed_comparators |
| if "allowed_operators" not in chain_kwargs: |
| chain_kwargs[ |
| "allowed_operators" |
| ] = structured_query_translator.allowed_operators |
| llm_chain = load_query_constructor_chain( |
| llm, document_contents, metadata_field_info, **chain_kwargs |
| ) |
| return cls( |
| llm_chain=llm_chain, |
| vectorstore=vectorstore, |
| structured_query_translator=structured_query_translator, |
| **kwargs, |
| ) |
|
|
