langchain-b12 0.1.1__tar.gz → 0.1.3__tar.gz

{langchain_b12-0.1.1 → langchain_b12-0.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langchain-b12
-Version: 0.1.1
+Version: 0.1.3
 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.1 → langchain_b12-0.1.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "langchain-b12"
-version = "0.1.1"
+version = "0.1.3"
 description = "A reusable collection of tools and implementations for Langchain"
 readme = "README.md"
 authors = [
@@ -17,6 +17,11 @@ google = [
 ]
 dev = [
     "pytest>=8.3.5",
+    "pytest-asyncio>=1.0.0",
+]
+citations = [
+    "fuzzysearch>=0.8.0",
+    "langgraph>=0.4.7",
 ]
 
 [build-system]
@@ -59,3 +64,8 @@ reportUnknownParameterType = false
 reportUnknownMemberType = false
 reportUnknownArgumentType = false
 
+# Add pytest configuration
+[tool.pytest.ini_options]
+asyncio_default_fixture_loop_scope = "function"
+asyncio_mode = "auto"
+

langchain_b12-0.1.3/src/langchain_b12/citations/citations.py

@@ -0,0 +1,274 @@
+import re
+from collections.abc import Sequence
+from typing import Literal, TypedDict
+
+from fuzzysearch import find_near_matches
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import AIMessage, BaseMessage, SystemMessage
+from langchain_core.runnables import Runnable
+from langgraph.utils.runnable import RunnableCallable
+from pydantic import BaseModel, Field
+
+SYSTEM_PROMPT = """
+You are an expert at identifying and adding citations to text.
+Your task is to identify, for each sentence in the final message, which citations were used to generate it.
+
+You will receive a numbered zero-indexed list of sentences in the final message, e.g.
+```
+0: Grass is green.
+1: The sky is blue and the sun is shining.
+```
+The rest of the conversation may contain contexts enclosed in xml tags, e.g.
+```
+<context key="abc">
+Today is a sunny day and the color of the grass is green.
+</context>
+```
+Each sentence may have zero, one, or multiple citations from the contexts.
+Each citation may be used for zero, one or multiple sentences.
+A context may be cited zero, one, or multiple times.
+
+The final message will be based on the contexts, but may not mention them explicitly.
+You must identify which contexts and which parts of the contexts were used to generate each sentence.
+For each such case, you must return a citation with a "sentence_index", "cited_text" and "key" property.
+The "sentence_index" is the index of the sentence in the final message.
+The "cited_text" must be a substring of the full context that was used to generate the sentence.
+The "key" must be the key of the context that was used to generate the sentence.
+Make sure that you copy the "cited_text" verbatim from the context, or it will not be considered valid.
+
+For the example above, the output should look like this:
+[
+    {
+        "sentence_index": 0,
+        "cited_text": "the color of the grass is green",
+        "key": "abc"
+    },
+    {
+        "sentence_index": 1,
+        "cited_text": "Today is a sunny day",
+        "key": "abc"
+    },
+]
+""".strip()  # noqa: E501
+
+
+class Match(TypedDict):
+    start: int
+    end: int
+    dist: int
+    matched: str
+
+
+class CitationType(TypedDict):
+
+    cited_text: str | None
+    generated_cited_text: str
+    key: str
+    dist: int | None
+
+
+class ContentType(TypedDict):
+
+    citations: list[CitationType] | None
+    text: str
+    type: Literal["text"]
+
+
+class Citation(BaseModel):
+
+    sentence_index: int = Field(
+        ...,
+        description="The index of the sentence from your answer "
+        "that this citation refers to.",
+    )
+    cited_text: str = Field(
+        ...,
+        description="The text that is cited from the document. "
+        "Make sure you cite it verbatim!",
+    )
+    key: str = Field(..., description="The key of the document you are citing.")
+
+
+class Citations(BaseModel):
+
+    values: list[Citation] = Field(..., description="List of citations")
+
+
+def split_into_sentences(text: str) -> list[str]:
+    """Split text into sentences on punctuation marks and newlines."""
+    if not text:
+        return [text]
+
+    # Split after punctuation followed by spaces, or on newlines
+    # Use capturing groups to preserve delimiters (spaces and newlines)
+    parts = re.split(r"((?<=[.!?])(?= +)|\n+)", text)
+
+    # Filter out empty strings that can result from splitting
+    return [part for part in parts if part]
+
+
+def contains_context_tags(text: str) -> bool:
+    """Check if the text contains context tags."""
+    return bool(re.search(r"<context\s+key=[^>]+>.*?</context>", text, re.DOTALL))
+
+
+def merge_citations(
+    sentences: list[str], citations: list[tuple[Citation, Match | None]]
+) -> list[ContentType]:
+    """Merge citations into sentences."""
+    content: list[ContentType] = []
+    for sentence_index, sentence in enumerate(sentences):
+        _citations: list[CitationType] = []
+        for citation, match in citations:
+            if citation.sentence_index == sentence_index:
+                if match is None:
+                    _citations.append(
+                        {
+                            "cited_text": None,
+                            "generated_cited_text": citation.cited_text,
+                            "key": citation.key,
+                            "dist": None,
+                        }
+                    )
+                else:
+                    _citations.append(
+                        {
+                            "cited_text": match["matched"],
+                            "generated_cited_text": citation.cited_text,
+                            "key": citation.key,
+                            "dist": match["dist"],
+                        }
+                    )
+        content.append(
+            {"text": sentence, "citations": _citations or None, "type": "text"}
+        )
+
+    return content
+
+
+def validate_citations(
+    citations: Citations,
+    messages: Sequence[BaseMessage],
+    sentences: list[str],
+) -> list[tuple[Citation, Match | None]]:
+    """Validate the citations. Invalid citations are dropped."""
+    n_sentences = len(sentences)
+
+    all_text = "\n".join(
+        str(msg.content) for msg in messages if isinstance(msg.content, str)
+    )
+
+    citations_with_matches: list[tuple[Citation, Match | None]] = []
+    for citation in citations.values:
+        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)
+        if not matches:
+            citations_with_matches.append((citation, None))
+        else:
+            match = matches[0]
+            citations_with_matches.append(
+                (
+                    citation,
+                    Match(
+                        start=match.start,
+                        end=match.end,
+                        dist=match.dist,
+                        matched=match.matched,
+                    ),
+                )
+            )
+    return citations_with_matches
+
+
+async def add_citations(
+    model: BaseChatModel,
+    messages: Sequence[BaseMessage],
+    message: AIMessage,
+    system_prompt: str,
+) -> AIMessage:
+    """Add citations to the message."""
+    if not message.content:
+        # Nothing to be done, for example in case of a tool call
+        return message
+
+    assert isinstance(
+        message.content, str
+    ), "Citation agent currently only supports string content."
+
+    if not contains_context_tags("\n".join(str(msg.content) for msg in messages)):
+        # No context tags, nothing to do
+        return message
+
+    sentences = split_into_sentences(message.content)
+
+    num_width = len(str(len(sentences)))
+    numbered_message = AIMessage(
+        content="\n".join(
+            f"{str(i).rjust(num_width)}: {sentence.strip()}"
+            for i, sentence in enumerate(sentences)
+        ),
+        name=message.name,
+    )
+    system_message = SystemMessage(system_prompt)
+    _messages = [system_message, *messages, numbered_message]
+
+    citations = await model.with_structured_output(Citations).ainvoke(_messages)
+    assert isinstance(
+        citations, Citations
+    ), f"Expected Citations from model invocation but got {type(citations)}"
+    citations = validate_citations(citations, messages, sentences)
+
+    message.content = merge_citations(sentences, citations)  # type: ignore[assignment]
+    return message
+
+
+def create_citation_model(
+    model: BaseChatModel,
+    citation_model: BaseChatModel | None = None,
+    system_prompt: str | None = None,
+) -> Runnable[Sequence[BaseMessage], AIMessage]:
+    """Take a base chat model and wrap it such that it adds citations to the messages.
+    Any contexts to be cited should be provided in the messages as XML tags,
+    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",
+        },
+    )
+
+    Args:
+        model: The base chat model to wrap.
+        citation_model: The model to use for extracting citations.
+            If None, the base model is used.
+        system_prompt: The system prompt to use for the citation model.
+            If None, a default prompt is used.
+    """
+    citation_model = citation_model or model
+    system_prompt = system_prompt or SYSTEM_PROMPT
+
+    async def ainvoke_with_citations(
+        messages: Sequence[BaseMessage],
+    ) -> AIMessage:
+        """Invoke the model and add citations to the AIMessage."""
+        ai_message = await model.ainvoke(messages)
+        assert isinstance(
+            ai_message, AIMessage
+        ), f"Expected AIMessage from model invocation but got {type(ai_message)}"
+        return await add_citations(citation_model, messages, ai_message, system_prompt)
+
+    return RunnableCallable(
+        func=None,  # TODO: Implement a sync version if needed
+        afunc=ainvoke_with_citations,
+    )

{langchain_b12-0.1.1 → langchain_b12-0.1.3}/src/langchain_b12/genai/genai_utils.py

@@ -24,22 +24,37 @@ def multi_content_to_part(
     Args:
         contents: A sequence of dictionaries representing content. Examples:
             [
-                {
+                { # Text content
                     "type": "text",
                     "text": "This is a text message"
                 },
-                {
+                { # Image content from base64 encoded string with OpenAI format
                     "type": "image_url",
                     "image_url": {
                         "url": f"data:{mime_type};base64,{encoded_artifact}"
                     },
                 },
-                {
+                { # Image content from base64 encoded string with LangChain format
+                    "type": "image",
+                    "source_type": "base64",
+                    "data": "<base64 string>",
+                    "mime_type": "image/jpeg",
+                },
+                { # Image content from URL
+                    "type": "image",
+                    "source_type": "url",
+                    "url": "https://...",
+                },
+                { # File content from base64 encoded string
+                    "type": "file",
+                    "source_type": "base64",
+                    "mime_type": "application/pdf",
+                    "data": "<base64 data string>",
+                },
+                { # File content from URL
                     "type": "file",
-                    "file": {
-                        "uri": f"gs://{bucket_name}/{file_name}",
-                        "mime_type": mime_type,
-                    }
+                    "source_type": "url",
+                    "url": "https://...",
                 }
             ]
     """
@@ -60,15 +75,51 @@ def multi_content_to_part(
             mime_type = header.split(":", 1)[1].split(";", 1)[0]
             data = base64.b64decode(encoded_data)
             parts.append(types.Part.from_bytes(data=data, mime_type=mime_type))
+        elif content["type"] == "image":
+            if "data" in content:
+                assert isinstance(content["data"], str), "Expected str data"
+                assert "mime_type" in content, "Expected 'mime_type' in content"
+                assert isinstance(content["mime_type"], str), "Expected str mime_type"
+                data = base64.b64decode(content["data"])
+                parts.append(
+                    types.Part.from_bytes(data=data, mime_type=content["mime_type"])
+                )
+            elif "url" in content:
+                assert isinstance(content["url"], str), "Expected str url"
+                mime_type = content.get("mime_type", None)
+                assert mime_type is None or isinstance(
+                    mime_type, str
+                ), "Expected str mime_type"
+                parts.append(
+                    types.Part.from_uri(file_uri=content["url"], mime_type=mime_type)
+                )
+            else:
+                raise ValueError(
+                    "Expected either 'data' or 'url' in content for image type"
+                )
         elif content["type"] == "file":
-            assert "file" in content, "Expected 'file' in content"
-            file = content["file"]
-            assert isinstance(file, dict), "Expected dict file"
-            assert "uri" in file, "Expected 'uri' in content['file']"
-            assert "mime_type" in file, "Expected 'mime_type' in content['file']"
-            parts.append(
-                types.Part.from_uri(file_uri=file["uri"], mime_type=file["mime_type"])
-            )
+            if "data" in content:
+                assert isinstance(content["data"], str), "Expected str data"
+                assert "mime_type" in content, "Expected 'mime_type' in content"
+                assert isinstance(content["mime_type"], str), "Expected str mime_type"
+                data = base64.b64decode(content["data"])
+                parts.append(
+                    types.Part.from_bytes(data=data, mime_type=content["mime_type"])
+                )
+            elif "url" in content:
+                assert isinstance(content["url"], str), "Expected str url"
+                assert content["url"], "File URI is required"
+                mime_type = content.get("mime_type", None)
+                assert mime_type is None or isinstance(
+                    mime_type, str
+                ), "Expected str mime_type"
+                parts.append(
+                    types.Part.from_uri(file_uri=content["url"], mime_type=mime_type)
+                )
+            else:
+                raise ValueError(
+                    "Expected either 'data' or 'url' in content for file type"
+                )
         else:
             raise ValueError(f"Unknown content type: {content['type']}")
     return parts