chatterer 0.1.6__tar.gz → 0.1.8__tar.gz

{chatterer-0.1.6 → chatterer-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: chatterer
-Version: 0.1.6
+Version: 0.1.8
 Summary: The highest-level interface for various LLM APIs.
 Requires-Python: >=3.12
 Description-Content-Type: text/markdown

{chatterer-0.1.6 → chatterer-0.1.8}/chatterer/__init__.py

@@ -1,4 +1,10 @@
 from .language_model import Chatterer
+from .messages import (
+    AIMessage,
+    BaseMessage,
+    HumanMessage,
+    SystemMessage,
+)
 from .strategies import (
     AoTPipeline,
     AoTPrompter,
@@ -7,6 +13,7 @@ from .strategies import (
 )
 from .tools import (
     anything_to_markdown,
+    citation_chunker,
     get_default_html_to_markdown_options,
     html_to_markdown,
     pdf_to_text,
@@ -24,4 +31,9 @@ __all__ = [
     "pdf_to_text",
     "get_default_html_to_markdown_options",
     "pyscripts_to_snippets",
+    "citation_chunker",
+    "BaseMessage",
+    "HumanMessage",
+    "SystemMessage",
+    "AIMessage",
 ]

{chatterer-0.1.6 → chatterer-0.1.8}/chatterer/language_model.py

@@ -8,15 +8,18 @@ from typing import (
     Type,
     TypeAlias,
     TypeVar,
+    cast,
+    overload,
 )
 
 from langchain_core.language_models.base import LanguageModelInput
 from langchain_core.language_models.chat_models import BaseChatModel
-from langchain_core.messages import HumanMessage
 from langchain_core.runnables.base import Runnable
 from langchain_core.runnables.config import RunnableConfig
 from pydantic import BaseModel, Field
 
+from .messages import AIMessage, BaseMessage, HumanMessage
+
 if TYPE_CHECKING:
     from instructor import Partial
 
@@ -32,19 +35,37 @@ class Chatterer(BaseModel):
     client: BaseChatModel
     structured_output_kwargs: dict[str, Any] = Field(default_factory=dict)
 
-    def __call__(self, messages: LanguageModelInput) -> str:
-        """
-        Generate text from the given input messages.
-
-        Args:
-            messages (LanguageModelInput): Input messages for the language model.
-                Can be one of the following types:
-                - str: A single string message.
-                - list[dict[str, str]]: A list of dictionaries with 'role' and 'content' keys.
-                - tuple[str, str]: A tuple of strings representing the role and content of a single message.
-                - list[BaseMessage]: A list of BaseMessage objects. (BaseMessage is a Pydantic model; e.g. can import AIMessage, HumanMessage, SystemMessage from langchain_core.messages)
-        """
-        return self.generate(messages)
+    @overload
+    def __call__(
+        self,
+        messages: LanguageModelInput,
+        response_model: Type[PydanticModelT],
+        config: Optional[RunnableConfig] = None,
+        stop: Optional[list[str]] = None,
+        **kwargs: Any,
+    ) -> PydanticModelT: ...
+
+    @overload
+    def __call__(
+        self,
+        messages: LanguageModelInput,
+        response_model: None = None,
+        config: Optional[RunnableConfig] = None,
+        stop: Optional[list[str]] = None,
+        **kwargs: Any,
+    ) -> str: ...
+
+    def __call__(
+        self,
+        messages: LanguageModelInput,
+        response_model: Optional[Type[PydanticModelT]] = None,
+        config: Optional[RunnableConfig] = None,
+        stop: Optional[list[str]] = None,
+        **kwargs: Any,
+    ) -> str | PydanticModelT:
+        if response_model:
+            return self.generate_pydantic(response_model, messages, config, stop, **kwargs)
+        return self.client.invoke(input=messages, config=config, stop=stop, **kwargs).text()
 
     @classmethod
     def openai(
@@ -233,6 +254,40 @@ class Chatterer(BaseModel):
             )
         ])
 
+    @staticmethod
+    def get_num_tokens_from_message(message: BaseMessage) -> Optional[tuple[int, int]]:
+        try:
+            if isinstance(message, AIMessage) and (usage_metadata := message.usage_metadata):
+                input_tokens = int(usage_metadata["input_tokens"])
+                output_tokens = int(usage_metadata["output_tokens"])
+            else:
+                # Dynamic extraction for unknown structures
+                input_tokens: Optional[int] = None
+                output_tokens: Optional[int] = None
+
+                def _find_tokens(obj: object) -> None:
+                    nonlocal input_tokens, output_tokens
+                    if isinstance(obj, dict):
+                        for key, value in cast(dict[object, object], obj).items():
+                            if isinstance(value, int):
+                                if "input" in str(key) or "prompt" in str(key):
+                                    input_tokens = value
+                                elif "output" in str(key) or "completion" in str(key):
+                                    output_tokens = value
+                            else:
+                                _find_tokens(value)
+                    elif isinstance(obj, list):
+                        for item in cast(list[object], obj):
+                            _find_tokens(item)
+
+                _find_tokens(message.model_dump())
+
+            if input_tokens is None or output_tokens is None:
+                return None
+            return input_tokens, output_tokens
+        except Exception:
+            return None
+
 
 def with_structured_output(
     client: BaseChatModel,
@@ -255,36 +310,25 @@ if __name__ == "__main__":
 
     # === Synchronous Tests ===
 
-    # 1. generate
+    # generate
     print("=== Synchronous generate ===")
-    result_sync = chatterer.generate(prompt)
+    result_sync = chatterer(prompt)
     print("Result (generate):", result_sync)
 
-    # 2. __call__
-    print("\n=== Synchronous __call__ ===")
-    result_call = chatterer(prompt)
-    print("Result (__call__):", result_call)
-
-    # 3. generate_stream
+    # generate_stream
     print("\n=== Synchronous generate_stream ===")
     for i, chunk in enumerate(chatterer.generate_stream(prompt)):
         print(f"Chunk {i}:", chunk)
 
-    # 4. generate_pydantic
+    # generate_pydantic
     print("\n=== Synchronous generate_pydantic ===")
-    try:
-        result_pydantic = chatterer.generate_pydantic(Propositions, prompt)
-        print("Result (generate_pydantic):", result_pydantic)
-    except Exception as e:
-        print("Error in generate_pydantic:", e)
+    result_pydantic = chatterer(prompt, Propositions)
+    print("Result (generate_pydantic):", result_pydantic)
 
-    # 5. generate_pydantic_stream
+    # generate_pydantic_stream
     print("\n=== Synchronous generate_pydantic_stream ===")
-    try:
-        for i, chunk in enumerate(chatterer.generate_pydantic_stream(Propositions, prompt)):
-            print(f"Pydantic Chunk {i}:", chunk)
-    except Exception as e:
-        print("Error in generate_pydantic_stream:", e)
+    for i, chunk in enumerate(chatterer.generate_pydantic_stream(Propositions, prompt)):
+        print(f"Pydantic Chunk {i}:", chunk)
 
     # === Asynchronous Tests ===
 

chatterer-0.1.8/chatterer/messages.py

@@ -0,0 +1,8 @@
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
+
+__all__ = [
+    "AIMessage",
+    "BaseMessage",
+    "HumanMessage",
+    "SystemMessage",
+]

{chatterer-0.1.6 → chatterer-0.1.8}/chatterer/strategies/atom_of_thoughts.py

@@ -6,11 +6,10 @@ from dataclasses import dataclass, field
 from enum import StrEnum
 from typing import Optional, Type, TypeVar
 
-from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
 from pydantic import BaseModel, Field, ValidationError
 
-# Import your Chatterer interface (do not remove)
 from ..language_model import Chatterer, LanguageModelInput
+from ..messages import AIMessage, BaseMessage, HumanMessage
 from .base import BaseStrategy
 
 # ---------------------------------------------------------------------------------

{chatterer-0.1.6 → chatterer-0.1.8}/chatterer/tools/__init__.py

@@ -1,3 +1,4 @@
+from .citation_chunking import citation_chunker
 from .convert_to_text import (
     anything_to_markdown,
     get_default_html_to_markdown_options,
@@ -12,4 +13,5 @@ __all__ = [
     "pdf_to_text",
     "get_default_html_to_markdown_options",
     "pyscripts_to_snippets",
+    "citation_chunker",
 ]

chatterer-0.1.8/chatterer/tools/citation_chunking/__init__.py

@@ -0,0 +1,3 @@
+from .citation_chunker import citation_chunker
+
+__all__ = ["citation_chunker"]

chatterer-0.1.8/chatterer/tools/citation_chunking/chunks.py

@@ -0,0 +1,53 @@
+import logging
+from typing import Callable, Optional, Self
+
+from pydantic import BaseModel, Field
+
+from ...language_model import Chatterer
+from ...messages import AIMessage, BaseMessage, HumanMessage
+from .prompt import (
+    generate_fewshot_affirmative_response,
+    generate_human_assistant_fewshot_examples,
+    generate_instruction,
+)
+from .reference import Reference
+
+logger = logging.getLogger(__name__)
+
+
+class CitationChunk(BaseModel):
+    subject: str = Field(description="The main topic or subject that the citations capture.")
+    references: list[Reference] = Field(description="A list of citation objects and/or regex patterns for the subject.")
+
+
+class CitationChunks(BaseModel):
+    citation_chunks: list[CitationChunk] = Field(
+        description="A list of citation chunks, each capturing a specific topic in the document."
+    )
+
+    @classmethod
+    def from_llm(
+        cls,
+        chatterer: Chatterer,
+        document: str,
+        fewshot_examples_generator: Optional[
+            Callable[[], list[tuple[str, str]]]
+        ] = generate_human_assistant_fewshot_examples,
+        instruction_generator: Optional[Callable[[], str]] = generate_instruction,
+        fewshot_affirmative_response: Optional[Callable[[], str]] = generate_fewshot_affirmative_response,
+    ) -> Self:
+        messages: list[BaseMessage] = []
+        if instruction_generator:
+            messages.append(HumanMessage(content=instruction_generator()))
+        if fewshot_examples_generator is not None:
+            if fewshot_affirmative_response:
+                messages.append(AIMessage(content=generate_fewshot_affirmative_response()))
+            for human_ask, ai_answer in fewshot_examples_generator():
+                messages.append(HumanMessage(content=human_ask))
+                messages.append(AIMessage(content=ai_answer))
+        messages.append(HumanMessage(content=document))
+        try:
+            return chatterer.generate_pydantic(response_model=cls, messages=messages)
+        except Exception as e:
+            logger.error(f"Error obtaining CitationChunks from LLM: {e}")
+            raise e

chatterer-0.1.8/chatterer/tools/citation_chunking/citation_chunker.py

@@ -0,0 +1,118 @@
+import logging
+from typing import Callable, NamedTuple, Optional, Self
+
+import colorama
+from colorama import Fore
+
+from ...language_model import Chatterer
+from .chunks import CitationChunks
+from .citations import Citations
+from .prompt import (
+    generate_fewshot_affirmative_response,
+    generate_human_assistant_fewshot_examples,
+    generate_instruction,
+)
+
+logger = logging.getLogger(__name__)
+colorama.init()
+
+
+class GlobalCoverage(NamedTuple):
+    coverage: float
+    matched_intervals: list[tuple[int, int]]
+
+    @staticmethod
+    def merge_intervals(intervals: list[tuple[int, int]]) -> list[tuple[int, int]]:
+        if not intervals:
+            return []
+        sorted_intervals = sorted(intervals, key=lambda x: x[0])
+        merged: list[tuple[int, int]] = [sorted_intervals[0]]
+        for current in sorted_intervals[1:]:
+            prev = merged[-1]
+            if current[0] <= prev[1]:
+                merged[-1] = (prev[0], max(prev[1], current[1]))
+            else:
+                merged.append(current)
+        return merged
+
+    @classmethod
+    def from_verified_citations(cls, verified_chunks: list[Citations], document: str) -> Self:
+        all_intervals: list[tuple[int, int]] = []
+        for chunk in verified_chunks:
+            for matches in chunk.references.values():
+                for m in matches:
+                    all_intervals.append((m.start_idx, m.end_idx))
+        merged: list[tuple[int, int]] = cls.merge_intervals(all_intervals)
+        doc_length: int = len(document)
+        total_matched = sum((e - s for s, e in merged))
+        coverage: float = total_matched / doc_length if doc_length > 0 else 0.0
+        return cls(coverage=coverage, matched_intervals=merged)
+
+
+def citation_chunker(
+    document: str,
+    chatterer: Chatterer,
+    global_coverage_threshold: float = 0.9,
+    num_refinement_steps: int = 3,
+    fewshot_examples_generator: Optional[
+        Callable[[], list[tuple[str, str]]]
+    ] = generate_human_assistant_fewshot_examples,
+    instruction_generator: Optional[Callable[[], str]] = generate_instruction,
+    fewshot_affirmative_response: Optional[Callable[[], str]] = generate_fewshot_affirmative_response,
+    test_global_coverage: bool = False,
+) -> list[Citations]:
+    """
+    1) Obtain CitationChunks via the LLM.
+    2) Process each chunk to extract MatchedText using snippet-based index correction.
+    3) Calculate overall document coverage and print results.
+    """
+    unverified_chunks: CitationChunks = CitationChunks.from_llm(
+        chatterer=chatterer,
+        document=document,
+        fewshot_examples_generator=fewshot_examples_generator,
+        instruction_generator=instruction_generator,
+        fewshot_affirmative_response=fewshot_affirmative_response,
+    )
+
+    verified_chunks: list[Citations] = []
+    for chunk in unverified_chunks.citation_chunks:
+        try:
+            vc: Citations = Citations.from_unverified(
+                unverified_chunk=chunk,
+                document=document,
+                model_and_refinement_steps=(chatterer, num_refinement_steps),
+            )
+            verified_chunks.append(vc)
+        except Exception as e:
+            logger.error(f"Error processing chunk for subject '{chunk.subject}': {e}")
+
+    if test_global_coverage:
+        gc = GlobalCoverage.from_verified_citations(verified_chunks, document)
+        logger.info(f"Global coverage: {gc.coverage * 100:.1f}%")
+        if gc.coverage < global_coverage_threshold:
+            logger.info(
+                f"Global coverage {gc.coverage * 100:.1f}% is below the threshold {global_coverage_threshold * 100:.1f}%."
+            )
+        print("=== Final Global Coverage Check ===")
+        print(f"Overall coverage: {gc.coverage * 100:.1f}% of the document.")
+        if gc.matched_intervals:
+            print("Merged matched intervals:")
+            for interval in gc.matched_intervals:
+                print(f" - {interval}")
+        else:
+            print("No matches found across all chunks.")
+        print("\n=== Raw Semantic Chunking Result ===")
+        for vc in verified_chunks:
+            print(f"{Fore.LIGHTGREEN_EX}[SUBJECT] {Fore.GREEN}{vc.name}{Fore.RESET}")
+            if vc.references:
+                for source_key, matches in vc.references.items():
+                    print(f"{Fore.LIGHTBLUE_EX}  [SOURCE] {Fore.BLUE}{source_key}{Fore.RESET}")
+                    for mt in matches:
+                        snippet = repr(mt.text)
+                        print(
+                            f"    {Fore.LIGHTYELLOW_EX}[MATCH @ {mt.start_idx}~{mt.end_idx}] {Fore.YELLOW}{snippet}{Fore.RESET}"
+                        )
+            else:
+                print(" - (No matches found even after refinement.)")
+
+    return verified_chunks