langchain-b12 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

langchain_b12/citations/citations.py

@@ -1,10 +1,12 @@
 import re
 from collections.abc import Sequence
-from typing import Literal, TypedDict
+from typing import Any, Literal, TypedDict
+from uuid import UUID
 
-from fuzzysearch import find_near_matches
+from langchain_core.callbacks import Callbacks
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage, BaseMessage, SystemMessage
+from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResult
 from langchain_core.runnables import Runnable
 from langgraph.utils.runnable import RunnableCallable
 from pydantic import BaseModel, Field
@@ -152,6 +154,8 @@ def validate_citations(
     sentences: list[str],
 ) -> list[tuple[Citation, Match | None]]:
     """Validate the citations. Invalid citations are dropped."""
+    from fuzzysearch import find_near_matches
+
     n_sentences = len(sentences)
 
     all_text = "\n".join(
@@ -163,7 +167,11 @@ def validate_citations(
         if citation.sentence_index < 0 or citation.sentence_index >= n_sentences:
             # discard citations that refer to non-existing sentences
             continue
-        matches = find_near_matches(citation.cited_text, all_text, max_l_dist=5)
+        # Allow for 10% error distance
+        max_l_dist = max(1, len(citation.cited_text) // 10)
+        matches = find_near_matches(
+            citation.cited_text, all_text, max_l_dist=max_l_dist
+        )
         if not matches:
             citations_with_matches.append((citation, None))
         else:
@@ -187,6 +195,7 @@ async def add_citations(
     messages: Sequence[BaseMessage],
     message: AIMessage,
     system_prompt: str,
+    **kwargs: Any,
 ) -> AIMessage:
     """Add citations to the message."""
     if not message.content:
@@ -214,7 +223,9 @@ async def add_citations(
     system_message = SystemMessage(system_prompt)
     _messages = [system_message, *messages, numbered_message]
 
-    citations = await model.with_structured_output(Citations).ainvoke(_messages)
+    citations = await model.with_structured_output(Citations).ainvoke(
+        _messages, **kwargs
+    )
     assert isinstance(
         citations, Citations
     ), f"Expected Citations from model invocation but got {type(citations)}"
@@ -234,18 +245,25 @@ def create_citation_model(
     e.g. `<context key="abc">Today is a sunny day</context>`.
     The returned AIMessage will have the following structure:
     AIMessage(
-        content= {
-            "citations": [
-                {
-                    "cited_text": "the color of the grass is green",
-                    "generated_cited_text": "the color of the grass is green",
-                    "key": "abc",
-                    "dist": 0,
-                }
-            ],
-            "text": "The grass is green",
-            "type": "text",
-        },
+        content=[
+            {
+                "citations": [
+                    {
+                        "cited_text": "the color of the grass is green",
+                        "generated_cited_text": "the color of the grass is green",
+                        "key": "abc",
+                        "dist": 0,
+                    }
+                ],
+                "text": "The grass is green",
+                "type": "text",
+            },
+            {
+                "citations": None,
+                "text": "Is there anything else I can help you with?",
+                "type": "text",
+            }
+        ]
     )
 
     Args:
@@ -272,3 +290,70 @@ def create_citation_model(
         func=None,  # TODO: Implement a sync version if needed
         afunc=ainvoke_with_citations,
     )
+
+
+class CitationMixin(BaseChatModel):
+    """Mixin class to add citation functionality to a runnable.
+
+    Example usage:
+    ```
+    from langchain_b12.genai.genai import ChatGenAI
+    from langchain_b12.citations.citations import CitationMixin
+
+    class CitationModel(ChatGenAI, CitationMixin):
+        pass
+    ```
+    """
+
+    async def agenerate(
+        self,
+        messages: list[list[BaseMessage]],
+        stop: list[str] | None = None,
+        callbacks: Callbacks = None,
+        *,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        run_name: str | None = None,
+        run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> LLMResult:
+        # Check if we should generate citations and remove it from kwargs
+        generate_citations = kwargs.pop("generate_citations", True)
+
+        llm_result = await super().agenerate(
+            messages,
+            stop,
+            callbacks,
+            tags=tags,
+            metadata=metadata,
+            run_name=run_name,
+            run_id=run_id,
+            **kwargs,
+        )
+
+        # Prevent recursion when extracting citations
+        if not generate_citations:
+            # Below we are call `add_citations` which will call `agenerate` again
+            # This will lead to an infinite loop if we don't stop here.
+            # We explicitly pass `generate_citations=False` below to sto this recursion.
+            return llm_result
+
+        # overwrite each generation with a version that has citations added
+        for _messages, generations in zip(messages, llm_result.generations):
+            for generation in generations:
+                assert isinstance(generation, ChatGeneration) and not isinstance(
+                    generation, ChatGenerationChunk
+                ), f"Expected ChatGeneration; received {type(generation)}"
+                assert isinstance(
+                    generation.message, AIMessage
+                ), f"Expected AIMessage; received {type(generation.message)}"
+                message_with_citations = await add_citations(
+                    self,
+                    _messages,
+                    generation.message,
+                    SYSTEM_PROMPT,
+                    generate_citations=False,
+                )
+                generation.message = message_with_citations
+
+        return llm_result

langchain_b12/genai/genai.py

@@ -90,6 +90,8 @@ class ChatGenAI(BaseChatModel):
                 HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
             }
             """  # noqa: E501
+    thinking_config: types.ThinkingConfig | None = None
+    "The thinking configuration to use for the model."
 
     model_config = ConfigDict(
         arbitrary_types_allowed=True,
@@ -208,6 +210,10 @@ class ChatGenAI(BaseChatModel):
                 candidate_count=self.n,
                 stop_sequences=stop or self.stop,
                 safety_settings=self.safety_settings,
+                thinking_config=self.thinking_config,
+                automatic_function_calling=types.AutomaticFunctionCallingConfig(
+                    disable=True,
+                ),
                 **kwargs,
             ),
         )
@@ -240,6 +246,10 @@ class ChatGenAI(BaseChatModel):
                 candidate_count=self.n,
                 stop_sequences=stop or self.stop,
                 safety_settings=self.safety_settings,
+                thinking_config=self.thinking_config,
+                automatic_function_calling=types.AutomaticFunctionCallingConfig(
+                    disable=True,
+                ),
                 **kwargs,
             ),
         )
@@ -362,6 +372,12 @@ class ChatGenAI(BaseChatModel):
                 input_tokens=usage_metadata.prompt_token_count or 0,
                 output_tokens=usage_metadata.candidates_token_count or 0,
                 total_tokens=usage_metadata.total_token_count or 0,
+                input_token_details={
+                    "cache_read": usage_metadata.cached_content_token_count or 0
+                },
+                output_token_details={
+                    "reasoning": usage_metadata.thoughts_token_count or 0
+                },
             )
 
         total_lc_usage: UsageMetadata | None = (

{langchain_b12-0.1.3.dist-info → langchain_b12-0.1.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langchain-b12
-Version: 0.1.3
+Version: 0.1.5
 Summary: A reusable collection of tools and implementations for Langchain
 Author-email: Vincent Min <vincent.min@b12-consulting.com>
 Requires-Python: >=3.11

langchain_b12-0.1.5.dist-info/RECORD

@@ -0,0 +1,9 @@
+langchain_b12/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+langchain_b12/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+langchain_b12/citations/citations.py,sha256=ZQvYayjQXIUaRosJ0qwL3Nc7kC8sBzmaIkE-BOslaVI,12261
+langchain_b12/genai/embeddings.py,sha256=od2bVIgt7v9aNAHG0PVypVF1H_XgHto2nTd8vwfvyN8,3355
+langchain_b12/genai/genai.py,sha256=7X7nDt76Icc5woV5b7FX_uza9YgFpFp1_PcYtXPriqE,16667
+langchain_b12/genai/genai_utils.py,sha256=tA6UiJURK25-11vtaX4768UV47jDCYwVKIIWydD4Egw,10736
+langchain_b12-0.1.5.dist-info/METADATA,sha256=unv3NxdFU_VrlPmIuTmDB2dHRi9go44B-q83kQgLUqI,1204
+langchain_b12-0.1.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+langchain_b12-0.1.5.dist-info/RECORD,,

langchain_b12-0.1.3.dist-info/RECORD

@@ -1,9 +0,0 @@
-langchain_b12/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-langchain_b12/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-langchain_b12/citations/citations.py,sha256=6HYKjyp6MaAWiLWZp-azQ5mM-drgt-Xytgarl7YwxhM,9321
-langchain_b12/genai/embeddings.py,sha256=od2bVIgt7v9aNAHG0PVypVF1H_XgHto2nTd8vwfvyN8,3355
-langchain_b12/genai/genai.py,sha256=gzkgtvs3wNjcslS_KFZYCajUZIsJkVN2Tq2Q1RMIPyc,15910
-langchain_b12/genai/genai_utils.py,sha256=tA6UiJURK25-11vtaX4768UV47jDCYwVKIIWydD4Egw,10736
-langchain_b12-0.1.3.dist-info/METADATA,sha256=gvKeYszVVVT37bI2RN8T3vOIafWAn48Pe9KTaDUeNd4,1204
-langchain_b12-0.1.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-langchain_b12-0.1.3.dist-info/RECORD,,