langchain-google-genai 0.0.10rc0__py3-none-any.whl → 1.0.1__py3-none-any.whl

langchain_google_genai/embeddings.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 
 # TODO: remove ignore once the google package is published with types
 import google.generativeai as genai  # type: ignore[import]
@@ -43,6 +43,13 @@ class GoogleGenerativeAIEmbeddings(BaseModel, Embeddings):
         description="The Google API key to use. If not provided, "
         "the GOOGLE_API_KEY environment variable will be used.",
     )
+    credentials: Any = Field(
+        default=None,
+        exclude=True,
+        description="The default custom credentials "
+        "(google.auth.credentials.Credentials) to use when making API calls. If not "
+        "provided, credentials will be ascertained from the GOOGLE_API_KEY envvar",
+    )
     client_options: Optional[Dict] = Field(
         None,
         description=(
@@ -58,17 +65,24 @@ class GoogleGenerativeAIEmbeddings(BaseModel, Embeddings):
     @root_validator()
     def validate_environment(cls, values: Dict) -> Dict:
         """Validates params and passes them to google-generativeai package."""
-        google_api_key = get_from_dict_or_env(
-            values, "google_api_key", "GOOGLE_API_KEY"
-        )
-        if isinstance(google_api_key, SecretStr):
-            google_api_key = google_api_key.get_secret_value()
-
-        genai.configure(
-            api_key=google_api_key,
-            transport=values.get("transport"),
-            client_options=values.get("client_options"),
-        )
+        if values.get("credentials"):
+            genai.configure(
+                credentials=values.get("credentials"),
+                transport=values.get("transport"),
+                client_options=values.get("client_options"),
+            )
+        else:
+            google_api_key = get_from_dict_or_env(
+                values, "google_api_key", "GOOGLE_API_KEY"
+            )
+            if isinstance(google_api_key, SecretStr):
+                google_api_key = google_api_key.get_secret_value()
+
+            genai.configure(
+                api_key=google_api_key,
+                transport=values.get("transport"),
+                client_options=values.get("client_options"),
+            )
         return values
 
     def _embed(

langchain_google_genai/genai_aqa.py

@@ -0,0 +1,134 @@
+"""Google GenerativeAI Attributed Question and Answering (AQA) service.
+
+The GenAI Semantic AQA API is a managed end to end service that allows
+developers to create responses grounded on specified passages based on
+a user query. For more information visit:
+https://developers.generativeai.google/guide
+"""
+
+from typing import Any, List, Optional
+
+import google.ai.generativelanguage as genai
+from langchain_core.pydantic_v1 import BaseModel, PrivateAttr
+from langchain_core.runnables import RunnableSerializable
+from langchain_core.runnables.config import RunnableConfig
+
+from . import _genai_extension as genaix
+
+
+class AqaInput(BaseModel):
+    """Input to `GenAIAqa.invoke`.
+
+    Attributes:
+        prompt: The user's inquiry.
+        source_passages: A list of passage that the LLM should use only to
+            answer the user's inquiry.
+    """
+
+    prompt: str
+    source_passages: List[str]
+
+
+class AqaOutput(BaseModel):
+    """Output from `GenAIAqa.invoke`.
+
+    Attributes:
+        answer: The answer to the user's inquiry.
+        attributed_passages: A list of passages that the LLM used to construct
+            the answer.
+        answerable_probability: The probability of the question being answered
+            from the provided passages.
+    """
+
+    answer: str
+    attributed_passages: List[str]
+    answerable_probability: float
+
+
+class _AqaModel(BaseModel):
+    """Wrapper for Google's internal AQA model."""
+
+    _client: genai.GenerativeServiceClient = PrivateAttr()
+    _answer_style: int = PrivateAttr()
+    _safety_settings: List[genai.SafetySetting] = PrivateAttr()
+    _temperature: Optional[float] = PrivateAttr()
+
+    def __init__(
+        self,
+        answer_style: int = genai.GenerateAnswerRequest.AnswerStyle.ABSTRACTIVE,
+        safety_settings: List[genai.SafetySetting] = [],
+        temperature: Optional[float] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._client = genaix.build_generative_service()
+        self._answer_style = answer_style
+        self._safety_settings = safety_settings
+        self._temperature = temperature
+
+    def generate_answer(
+        self,
+        prompt: str,
+        passages: List[str],
+    ) -> genaix.GroundedAnswer:
+        return genaix.generate_answer(
+            prompt=prompt,
+            passages=passages,
+            client=self._client,
+            answer_style=self._answer_style,
+            safety_settings=self._safety_settings,
+            temperature=self._temperature,
+        )
+
+
+class GenAIAqa(RunnableSerializable[AqaInput, AqaOutput]):
+    """Google's Attributed Question and Answering service.
+
+    Given a user's query and a list of passages, Google's server will return
+    a response that is grounded to the provided list of passages. It will not
+    base the response on parametric memory.
+
+    Attributes:
+        answer_style: keyword-only argument. See
+            `google.ai.generativelanguage.AnswerStyle` for details.
+    """
+
+    # Actual type is .aqa_model.AqaModel.
+    _client: _AqaModel = PrivateAttr()
+
+    # Actual type is genai.AnswerStyle.
+    # 1 = ABSTRACTIVE.
+    # Cannot use the actual type here because user may not have
+    # google.generativeai installed.
+    answer_style: int = 1
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Construct a Google Generative AI AQA model.
+
+        All arguments are optional.
+
+        Args:
+            answer_style: See
+              `google.ai.generativelanguage.GenerateAnswerRequest.AnswerStyle`.
+            safety_settings: See `google.ai.generativelanguage.SafetySetting`.
+            temperature: 0.0 to 1.0.
+        """
+        super().__init__(**kwargs)
+        self._client = _AqaModel(**kwargs)
+
+    def invoke(
+        self, input: AqaInput, config: Optional[RunnableConfig] = None
+    ) -> AqaOutput:
+        """Generates a grounded response using the provided passages."""
+
+        response = self._client.generate_answer(
+            prompt=input.prompt, passages=input.source_passages
+        )
+
+        return AqaOutput(
+            answer=response.answer,
+            attributed_passages=[
+                passage.text for passage in response.attributed_passages
+            ],
+            answerable_probability=response.answerable_probability or 0.0,
+        )

langchain_google_genai/google_vector_store.py

@@ -0,0 +1,493 @@
+"""Google Generative AI Vector Store.
+
+The GenAI Semantic Retriever API is a managed end-to-end service that allows
+developers to create a corpus of documents to perform semantic search on
+related passages given a user query. For more information visit:
+https://developers.generativeai.google/guide
+"""
+
+import asyncio
+from functools import partial
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Tuple,
+)
+
+import google.ai.generativelanguage as genai
+from langchain_core.documents import Document
+from langchain_core.embeddings import Embeddings
+from langchain_core.pydantic_v1 import BaseModel, PrivateAttr
+from langchain_core.runnables import Runnable, RunnableLambda, RunnablePassthrough
+from langchain_core.vectorstores import VectorStore
+
+from . import _genai_extension as genaix
+from .genai_aqa import (
+    AqaInput,
+    AqaOutput,
+    GenAIAqa,
+)
+
+
+class ServerSideEmbedding(Embeddings):
+    """Do nothing embedding model where the embedding is done by the server."""
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        return [[] for _ in texts]
+
+    def embed_query(self, text: str) -> List[float]:
+        return []
+
+
+class DoesNotExistsException(Exception):
+    def __init__(self, *, corpus_id: str, document_id: Optional[str] = None) -> None:
+        if document_id is None:
+            message = f"No such corpus {corpus_id}"
+        else:
+            message = f"No such document {document_id} under corpus {corpus_id}"
+        super().__init__(message)
+
+
+class _SemanticRetriever(BaseModel):
+    """Wrapper class to Google's internal semantric retriever service."""
+
+    name: genaix.EntityName
+    _client: genai.RetrieverServiceClient = PrivateAttr()
+
+    def __init__(self, *, client: genai.RetrieverServiceClient, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._client = client
+
+    @classmethod
+    def from_ids(
+        cls, corpus_id: str, document_id: Optional[str]
+    ) -> "_SemanticRetriever":
+        name = genaix.EntityName(corpus_id=corpus_id, document_id=document_id)
+        client = genaix.build_semantic_retriever()
+
+        # Check the entity exists on Google server.
+        if name.is_corpus():
+            if genaix.get_corpus(corpus_id=corpus_id, client=client) is None:
+                raise DoesNotExistsException(corpus_id=corpus_id)
+        elif name.is_document():
+            assert document_id is not None
+            if (
+                genaix.get_document(
+                    corpus_id=corpus_id, document_id=document_id, client=client
+                )
+                is None
+            ):
+                raise DoesNotExistsException(
+                    corpus_id=corpus_id, document_id=document_id
+                )
+
+        return cls(name=name, client=client)
+
+    def add_texts(
+        self,
+        texts: Iterable[str],
+        metadatas: Optional[List[Dict[str, Any]]] = None,
+        document_id: Optional[str] = None,
+    ) -> List[str]:
+        if self.name.document_id is None and document_id is None:
+            raise NotImplementedError(
+                "Adding texts to a corpus directly is not supported. "
+                "Please provide a document ID under the corpus first. "
+                "Then add the texts to the document."
+            )
+        if (
+            self.name.document_id is not None
+            and document_id is not None
+            and self.name.document_id != document_id
+        ):
+            raise NotImplementedError(
+                f"Parameter `document_id` {document_id} does not match the "
+                f"vector store's `document_id` {self.name.document_id}"
+            )
+        assert self.name.document_id or document_id is not None
+        new_document_id = self.name.document_id or document_id or ""
+
+        texts = list(texts)
+        if metadatas is None:
+            metadatas = [{} for _ in texts]
+        if len(texts) != len(metadatas):
+            raise ValueError(
+                f"metadatas's length {len(metadatas)} and "
+                f"texts's length {len(texts)} are mismatched"
+            )
+
+        chunks = genaix.batch_create_chunk(
+            corpus_id=self.name.corpus_id,
+            document_id=new_document_id,
+            texts=texts,
+            metadatas=metadatas,
+            client=self._client,
+        )
+
+        return [chunk.name for chunk in chunks if chunk.name]
+
+    def similarity_search(
+        self,
+        query: str,
+        k: int = 4,
+        filter: Optional[Dict[str, Any]] = None,
+    ) -> List[Tuple[str, float]]:
+        if self.name.is_corpus():
+            relevant_chunks = genaix.query_corpus(
+                corpus_id=self.name.corpus_id,
+                query=query,
+                k=k,
+                filter=filter,
+                client=self._client,
+            )
+        else:
+            assert self.name.is_document()
+            assert self.name.document_id is not None
+            relevant_chunks = genaix.query_document(
+                corpus_id=self.name.corpus_id,
+                document_id=self.name.document_id,
+                query=query,
+                k=k,
+                filter=filter,
+                client=self._client,
+            )
+
+        return [
+            (chunk.chunk.data.string_value, chunk.chunk_relevance_score)
+            for chunk in relevant_chunks
+        ]
+
+    def delete(self, ids: Optional[List[str]] = None, **kwargs: Any) -> Optional[bool]:
+        for id in ids or []:
+            name = genaix.EntityName.from_str(id)
+            _delete_chunk(
+                corpus_id=name.corpus_id,
+                document_id=name.document_id,
+                chunk_id=name.chunk_id,
+                client=self._client,
+            )
+        return True
+
+
+def _delete_chunk(
+    *,
+    corpus_id: str,
+    document_id: Optional[str],
+    chunk_id: Optional[str],
+    client: genai.RetrieverServiceClient,
+) -> None:
+    if chunk_id is not None:
+        if document_id is None:
+            raise ValueError(f"Chunk {chunk_id} requires a document ID")
+        genaix.delete_chunk(
+            corpus_id=corpus_id,
+            document_id=document_id,
+            chunk_id=chunk_id,
+            client=client,
+        )
+    elif document_id is not None:
+        genaix.delete_document(
+            corpus_id=corpus_id, document_id=document_id, client=client
+        )
+    else:
+        genaix.delete_corpus(corpus_id=corpus_id, client=client)
+
+
+class GoogleVectorStore(VectorStore):
+    """Google GenerativeAI Vector Store.
+
+    Currently, it computes the embedding vectors on the server side.
+
+    Example: Add texts to an existing corpus.
+
+        store = GoogleVectorStore(corpus_id="123")
+        store.add_documents(documents, document_id="456")
+
+    Example: Create a new corpus.
+
+        store = GoogleVectorStore.create_corpus(
+            corpus_id="123", display_name="My Google corpus")
+
+    Example: Query the corpus for relevant passages.
+
+        store.as_retriever() \
+            .get_relevant_documents("Who caught the gingerbread man?")
+
+    Example: Ask the corpus for grounded responses!
+
+        aqa = store.as_aqa()
+        response = aqa.invoke("Who caught the gingerbread man?")
+        print(response.answer)
+        print(response.attributed_passages)
+        print(response.answerability_probability)
+
+    You can also operate at Google's Document level.
+
+    Example: Add texts to an existing Google Vector Store Document.
+
+        doc_store = GoogleVectorStore(corpus_id="123", document_id="456")
+        doc_store.add_documents(documents)
+
+    Example: Create a new Google Vector Store Document.
+
+        doc_store = GoogleVectorStore.create_document(
+            corpus_id="123", document_id="456", display_name="My Google document")
+
+    Example: Query the Google document.
+
+        doc_store.as_retriever() \
+            .get_relevant_documents("Who caught the gingerbread man?")
+
+    For more details, see the class's methods.
+    """
+
+    _retriever: _SemanticRetriever
+
+    def __init__(
+        self, *, corpus_id: str, document_id: Optional[str] = None, **kwargs: Any
+    ):
+        """Returns an existing Google Semantic Retriever corpus or document.
+
+        If just the corpus ID is provided, the vector store operates over all
+        documents within that corpus.
+
+        If the document ID is provided, the vector store operates over just that
+        document.
+
+        Raises:
+            DoesNotExistsException if the IDs do not match to anything on Google
+                server. In this case, consider using `create_corpus` or
+                `create_document` to create one.
+        """
+        super().__init__(**kwargs)
+        self._retriever = _SemanticRetriever.from_ids(corpus_id, document_id)
+
+    @classmethod
+    def create_corpus(
+        cls,
+        corpus_id: Optional[str] = None,
+        display_name: Optional[str] = None,
+    ) -> "GoogleVectorStore":
+        """Create a Google Semantic Retriever corpus.
+
+        Args:
+            corpus_id: The ID to use to create the new corpus. If not provided,
+                Google server will provide one.
+            display_name: The title of the new corpus. If not provided, Google
+                server will provide one.
+
+        Returns:
+            An instance of vector store that points to the newly created corpus.
+        """
+        client = genaix.build_semantic_retriever()
+        corpus = genaix.create_corpus(
+            corpus_id=corpus_id, display_name=display_name, client=client
+        )
+
+        n = genaix.EntityName.from_str(corpus.name)
+        return cls(corpus_id=n.corpus_id)
+
+    @classmethod
+    def create_document(
+        cls,
+        corpus_id: str,
+        document_id: Optional[str] = None,
+        display_name: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> "GoogleVectorStore":
+        """Create a Google Semantic Retriever document.
+
+        Args:
+            corpus_id: ID of an existing corpus.
+            document_id: The ID to use to create the new Google Semantic
+                Retriever document. If not provided, Google server will provide
+                one.
+            display_name: The title of the new document. If not provided, Google
+                server will provide one.
+
+        Returns:
+            An instance of vector store that points to the newly created
+            document.
+        """
+        client = genaix.build_semantic_retriever()
+        document = genaix.create_document(
+            corpus_id=corpus_id,
+            document_id=document_id,
+            display_name=display_name,
+            metadata=metadata,
+            client=client,
+        )
+
+        assert document.name is not None
+        d = genaix.EntityName.from_str(document.name)
+        return cls(corpus_id=d.corpus_id, document_id=d.document_id)
+
+    @classmethod
+    def from_texts(
+        cls,
+        texts: List[str],
+        embedding: Optional[Embeddings] = None,
+        metadatas: Optional[List[dict[str, Any]]] = None,
+        *,
+        corpus_id: Optional[str] = None,  # str required
+        document_id: Optional[str] = None,  # str required
+        **kwargs: Any,
+    ) -> "GoogleVectorStore":
+        """Returns a vector store of an existing document with the specified text.
+
+        Args:
+            corpus_id: REQUIRED. Must be an existing corpus.
+            document_id: REQUIRED. Must be an existing document.
+            texts: Texts to be loaded into the vector store.
+
+        Returns:
+            A vector store pointing to the specified Google Semantic Retriever
+            Document.
+
+        Raises:
+            DoesNotExistsException if the IDs do not match to anything at
+                Google server.
+        """
+        if corpus_id is None or document_id is None:
+            raise NotImplementedError(
+                "Must provide an existing corpus ID and document ID"
+            )
+
+        doc_store = cls(corpus_id=corpus_id, document_id=document_id, **kwargs)
+        doc_store.add_texts(texts, metadatas)
+
+        return doc_store
+
+    @property
+    def name(self) -> str:
+        """Returns the name of the Google entity.
+
+        You shouldn't need to care about this unless you want to access your
+        corpus or document via Google Generative AI API.
+        """
+        return str(self._retriever.name)
+
+    @property
+    def corpus_id(self) -> str:
+        """Returns the corpus ID managed by this vector store."""
+        return self._retriever.name.corpus_id
+
+    @property
+    def document_id(self) -> Optional[str]:
+        """Returns the document ID managed by this vector store."""
+        return self._retriever.name.document_id
+
+    def add_texts(
+        self,
+        texts: Iterable[str],
+        metadatas: Optional[List[Dict[str, Any]]] = None,
+        *,
+        document_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> List[str]:
+        """Add texts to the vector store.
+
+        If the vector store points to a corpus (instead of a document), you must
+        also provide a `document_id`.
+
+        Returns:
+            Chunk's names created on Google servers.
+        """
+        return self._retriever.add_texts(texts, metadatas, document_id)
+
+    def similarity_search(
+        self,
+        query: str,
+        k: int = 4,
+        filter: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> List[Document]:
+        """Search the vector store for relevant texts."""
+        return [
+            document
+            for document, _ in self.similarity_search_with_score(
+                query, k, filter, **kwargs
+            )
+        ]
+
+    def similarity_search_with_score(
+        self,
+        query: str,
+        k: int = 4,
+        filter: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> List[Tuple[Document, float]]:
+        """Run similarity search with distance."""
+        return [
+            (Document(page_content=text), score)
+            for text, score in self._retriever.similarity_search(query, k, filter)
+        ]
+
+    def delete(self, ids: Optional[List[str]] = None, **kwargs: Any) -> Optional[bool]:
+        """Delete chunnks.
+
+        Note that the "ids" are not corpus ID or document ID. Rather, these
+        are the entity names returned by `add_texts`.
+
+        Returns:
+            True if successful. Otherwise, you should get an exception anyway.
+        """
+        return self._retriever.delete(ids)
+
+    async def adelete(
+        self, ids: Optional[List[str]] = None, **kwargs: Any
+    ) -> Optional[bool]:
+        return await asyncio.get_running_loop().run_in_executor(
+            None, partial(self.delete, **kwargs), ids
+        )
+
+    def _select_relevance_score_fn(self) -> Callable[[float], float]:
+        """
+        TODO: Check with the team about this!
+        The underlying vector store already returns a "score proper",
+        i.e. one in [0, 1] where higher means more *similar*.
+        """
+        return lambda score: score
+
+    def as_aqa(self, **kwargs: Any) -> Runnable[str, AqaOutput]:
+        """Construct a Google Generative AI AQA engine.
+
+        All arguments are optional.
+
+        Args:
+            answer_style: See
+              `google.ai.generativelanguage.GenerateAnswerRequest.AnswerStyle`.
+            safety_settings: See `google.ai.generativelanguage.SafetySetting`.
+            temperature: 0.0 to 1.0.
+        """
+        return (
+            RunnablePassthrough[str]()
+            | {
+                "prompt": RunnablePassthrough(),
+                "passages": self.as_retriever(),
+            }
+            | RunnableLambda(_toAqaInput)
+            | GenAIAqa(**kwargs)
+        )
+
+
+def _toAqaInput(input: Dict[str, Any]) -> AqaInput:
+    prompt = input["prompt"]
+    assert isinstance(prompt, str)
+
+    passages = input["passages"]
+    assert isinstance(passages, list)
+
+    source_passages: List[str] = []
+    for passage in passages:
+        assert isinstance(passage, Document)
+        source_passages.append(passage.page_content)
+
+    return AqaInput(
+        prompt=prompt,
+        source_passages=source_passages,
+    )