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

langchain_b12/citations/citations.py

@@ -1,10 +1,13 @@
 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
@@ -163,7 +166,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 +194,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 +222,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)}"
@@ -272,3 +282,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-0.1.3.dist-info → langchain_b12-0.1.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langchain-b12
-Version: 0.1.3
+Version: 0.1.4
 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.3.dist-info → langchain_b12-0.1.4.dist-info}/RECORD

@@ -1,9 +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=6HYKjyp6MaAWiLWZp-azQ5mM-drgt-Xytgarl7YwxhM,9321
+langchain_b12/citations/citations.py,sha256=FO9ErybQws082JvV-MtTj81fVzdUWxrALcS81ElRsMw,12023
 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,,
+langchain_b12-0.1.4.dist-info/METADATA,sha256=x659l7J9-4XSfjgZgGvR-cVoiCtTqq7cIaujV9JsTrE,1204
+langchain_b12-0.1.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+langchain_b12-0.1.4.dist-info/RECORD,,