autochatlib 0.1.0__py3-none-any.whl

autochat/__init__.py

@@ -0,0 +1,66 @@
+from .chat import AutoChat
+from .compression import (
+    AutoCompress,
+    CompressionStrategy,
+    KeepLatestN,
+    SummarizeAll,
+    SummarizeLatestN,
+)
+from .config import ChatConfig
+from .exceptions import (
+    AutoChatError,
+    AutoChatToolError,
+)
+from .guidelines import ChatGuideline
+from .retrieval import (
+    ChatRetriever,
+    RAGStrategy,
+    RetrievedDocument,
+    RetrievalConfig,
+    RetrievalInvocation,
+    RetrievalResult,
+    RetrieverFn,
+    RetrieverPostprocessor,
+    RetrieverPreprocessor,
+    SimpleRAGStrategy,
+)
+from .runtime import (
+    ChatRuntime,
+)
+from .tools import (
+    ChatTool,
+    ToolInvocation,
+    ToolPostprocessor,
+    ToolPreprocessor,
+    chat_tool,
+)
+
+__version__ = "0.1.0"
+__all__ = [
+    "AutoChat",
+    "AutoCompress",
+    "ChatConfig",
+    "CompressionStrategy",
+    "KeepLatestN",
+    "SummarizeAll",
+    "SummarizeLatestN",
+    "AutoChatError",
+    "AutoChatToolError",
+    "ChatRuntime",
+    "ChatRetriever",
+    "RetrievalConfig",
+    "RetrievedDocument",
+    "RetrievalResult",
+    "RAGStrategy",
+    "SimpleRAGStrategy",
+    "RetrievalInvocation",
+    "RetrieverFn",
+    "RetrieverPostprocessor",
+    "RetrieverPreprocessor",
+    "ChatTool",
+    "ToolInvocation",
+    "ToolPostprocessor",
+    "ToolPreprocessor",
+    "chat_tool",
+    "ChatGuideline",
+]

autochat/chat.py

@@ -0,0 +1,116 @@
+from typing import Any, AsyncIterator, Generic, Literal, Mapping, Sequence, TypeVar
+from uuid import uuid4
+
+from langchain_core.messages import HumanMessage
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.base import BaseCheckpointSaver
+
+from autochat.config import ChatConfig
+from autochat.compression import AutoCompress
+from autochat.graph.builder import build_chat_graph
+from autochat.graph.runtime import with_runtime_config
+from autochat.guidelines import ChatGuideline
+from autochat.retrieval import ChatRetriever, RetrievalConfig
+from autochat.runtime import ChatRuntime
+from autochat.tools import ChatTool
+
+TContext = TypeVar("TContext")
+
+
+class AutoChat(Generic[TContext]):
+    def __init__(
+        self,
+        *,
+        config: ChatConfig,
+        tools: Sequence[ChatTool[TContext, Any, Any]] = (),
+        retrievers: Sequence[ChatRetriever[TContext]] = (),
+        retrieval: RetrievalConfig[TContext] | None = None,
+        persistence: BaseCheckpointSaver | None = None,
+        compression: AutoCompress[TContext] | None = None,
+        system_message: str | None = None,
+        guidelines: Sequence[ChatGuideline] = (),
+    ) -> None:
+        self.config = config
+        self.tools = tuple(tools)
+        self.retrievers = tuple(retrievers)
+        self.retrieval = retrieval or RetrievalConfig()
+        self.system_message = system_message
+        self.guidelines = tuple(guidelines)
+        self.persistence = persistence
+        self.compression = compression
+
+        self._graph = build_chat_graph(
+            config=config,
+            tools=self.tools,
+            retrievers=self.retrievers,
+            retrieval=self.retrieval,
+            compression=compression,
+            system_message=system_message,
+            guidelines=self.guidelines,
+            persistence=persistence,
+        )
+
+    def _runtime(
+        self,
+        *,
+        thread_id: str,
+        context: TContext,
+        run_id: str | None,
+        metadata: Mapping[str, Any] | None,
+    ) -> ChatRuntime[TContext]:
+        return ChatRuntime(
+            thread_id=thread_id,
+            context=context,
+            run_id=run_id or str(uuid4()),
+            metadata=metadata or {},
+        )
+
+    async def ainvoke(
+        self,
+        input: str,
+        *,
+        thread_id: str,
+        context: TContext,
+        run_id: str | None = None,
+        metadata: Mapping[str, Any] | None = None,
+        config: RunnableConfig | None = None,
+    ) -> Any:
+        runtime = self._runtime(
+            thread_id=thread_id,
+            context=context,
+            run_id=run_id,
+            metadata=metadata,
+        )
+
+        graph_config = with_runtime_config(config, runtime)
+
+        return await self._graph.ainvoke(
+            {"messages": [HumanMessage(content=input)]}, config=graph_config
+        )
+
+    async def astream_events(
+        self,
+        input: str,
+        *,
+        thread_id: str,
+        context: TContext,
+        run_id: str | None = None,
+        metadata: Mapping[str, Any] | None = None,
+        config: RunnableConfig | None = None,
+        version: Literal["v1", "v2"] = "v2",
+    ) -> AsyncIterator[Any]:
+        runtime = self._runtime(
+            thread_id=thread_id,
+            context=context,
+            run_id=run_id,
+            metadata=metadata,
+        )
+
+        graph_config = with_runtime_config(config, runtime)
+
+        async for event in self._graph.astream_events(
+            {"messages": [HumanMessage(content=input)]},
+            config=graph_config,
+            version=version,
+        ):
+            yield event

autochat/compression/__init__.py

@@ -0,0 +1,11 @@
+from .config import AutoCompress
+from .strategies import KeepLatestN, SummarizeAll, SummarizeLatestN
+from .types import CompressionStrategy
+
+__all__ = [
+    "AutoCompress",
+    "CompressionStrategy",
+    "KeepLatestN",
+    "SummarizeAll",
+    "SummarizeLatestN",
+]

autochat/compression/config.py

@@ -0,0 +1,86 @@
+import warnings
+from dataclasses import dataclass, field
+from typing import Generic, TypeVar, cast
+
+from langchain.chat_models import BaseChatModel
+from langchain_core.messages import BaseMessage
+
+from autochat.config import ChatConfig
+from autochat.runtime import ChatRuntime
+
+from .strategies import SummarizeAll
+from .types import CompressionStrategy
+
+TContext = TypeVar("TContext")
+
+
+def _message_content_text(message: BaseMessage) -> str:
+    content = message.content
+    if isinstance(content, str):
+        return content
+    return str(content)
+
+
+# TODO: Check if this approach is good
+def _approximate_message_tokens(messages: list[BaseMessage]) -> int:
+    # Rough fallback used when the model cannot count tokens without extra deps.
+    return max(1, sum(len(_message_content_text(message)) for message in messages) // 4)
+
+
+def _count_message_tokens(
+    model: BaseChatModel,
+    messages: list[BaseMessage],
+) -> int:
+    try:
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                "ignore",
+                message="Using fallback GPT-2 tokenizer.*",
+                category=UserWarning,
+            )
+            return model.get_num_tokens_from_messages(messages)
+    except ImportError:
+        return _approximate_message_tokens(messages)
+
+
+@dataclass(frozen=True, slots=True)
+class AutoCompress(Generic[TContext]):
+    """Automatic chat-history compression policy."""
+
+    at: float = 0.6
+    strategy: CompressionStrategy[TContext] = field(
+        default_factory=lambda: cast(CompressionStrategy[TContext], SummarizeAll())
+    )
+    context_window: int | None = None
+    model: BaseChatModel | None = None
+
+    def __post_init__(self) -> None:
+        if not 0 < self.at <= 1:
+            raise ValueError(
+                "AutoCompress.at must be greater than 0 and less than or equal to 1."
+            )
+
+    async def acompress_if_needed(
+        self,
+        messages: list[BaseMessage],
+        *,
+        runtime: ChatRuntime[TContext],
+        config: ChatConfig,
+    ) -> list[BaseMessage] | None:
+        context_window = self.context_window or config.context_window
+        if context_window is None:
+            raise ValueError(
+                "AutoCompress requires ChatConfig.context_window or "
+                "AutoCompress.context_window."
+            )
+
+        model = self.model or config.model
+        used_tokens = _count_message_tokens(model, messages)
+        if used_tokens < context_window * self.at:
+            return None
+
+        return await self.strategy.acompress(
+            messages,
+            runtime=runtime,
+            model=model,
+        )

autochat/compression/strategies.py

@@ -0,0 +1,186 @@
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+from langchain.chat_models import BaseChatModel
+from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
+from langgraph.graph.message import REMOVE_ALL_MESSAGES, RemoveMessage
+
+from autochat.runtime import ChatRuntime
+
+TContext = TypeVar("TContext")
+
+_SUMMARY_METADATA = {
+    "autochat": {
+        "kind": "compression_summary",
+    }
+}
+
+
+def _summary_message(summary: str) -> SystemMessage:
+    return SystemMessage(
+        content=f"Previous conversation summary:\n{summary}",
+        additional_kwargs=_SUMMARY_METADATA,
+    )
+
+
+def _replace_messages(
+    messages: Sequence[BaseMessage],
+) -> list[BaseMessage]:
+    return [
+        RemoveMessage(id=REMOVE_ALL_MESSAGES),
+        *messages,
+    ]
+
+
+def _replace_with_summary_between(
+    messages_before: Sequence[BaseMessage],
+    summary: str,
+    messages_after: Sequence[BaseMessage],
+) -> list[BaseMessage]:
+    return _replace_messages(
+        [
+            *messages_before,
+            _summary_message(summary),
+            *messages_after,
+        ]
+    )
+
+
+def _replace_with_latest(
+    messages_to_keep: Sequence[BaseMessage],
+) -> list[BaseMessage]:
+    return _replace_messages(messages_to_keep)
+
+
+def _split_preserved_tail(
+    messages: Sequence[BaseMessage],
+    preserve_latest: int,
+) -> tuple[list[BaseMessage], list[BaseMessage]]:
+    message_list = list(messages)
+    if preserve_latest <= 0:
+        return message_list, []
+    return message_list[:-preserve_latest], message_list[-preserve_latest:]
+
+
+def _message_text(message: BaseMessage) -> str:
+    content = message.content
+    if isinstance(content, str):
+        return content
+    return str(content)
+
+
+def _format_messages_for_summary(messages: Sequence[BaseMessage]) -> str:
+    lines: list[str] = []
+    for message in messages:
+        role = message.type
+        name = getattr(message, "name", None)
+        label = f"{role}:{name}" if name else role
+        lines.append(f"{label}: {_message_text(message)}")
+    return "\n".join(lines)
+
+
+async def _summarize_messages(
+    messages: Sequence[BaseMessage],
+    *,
+    model: BaseChatModel,
+) -> str:
+    prompt = (
+        "Summarize this conversation for future continuation. "
+        "Preserve user goals, decisions, constraints, tool results, and unresolved tasks. "
+        "Do not add facts that are not present.\n\n"
+        f"Conversation:\n{_format_messages_for_summary(messages)}"
+    )
+    response = await model.ainvoke([HumanMessage(content=prompt)])
+    content = response.content
+    if isinstance(content, str):
+        return content
+    return str(content)
+
+
+@dataclass(frozen=True, slots=True)
+class KeepLatestN(Generic[TContext]):
+    """Drop older history and keep only the latest N messages."""
+
+    n: int
+
+    def __post_init__(self) -> None:
+        if self.n < 1:
+            raise ValueError("KeepLatestN.n must be at least 1.")
+
+    async def acompress(
+        self,
+        messages: Sequence[BaseMessage],
+        *,
+        runtime: ChatRuntime[TContext],
+        model: BaseChatModel,
+    ) -> list[BaseMessage]:
+        del runtime, model
+        return _replace_with_latest(list(messages)[-self.n :])
+
+
+@dataclass(frozen=True, slots=True)
+class SummarizeAll(Generic[TContext]):
+    """Summarize history into one summary and preserve the active latest messages."""
+
+    preserve_latest: int = 1
+
+    def __post_init__(self) -> None:
+        if self.preserve_latest < 0:
+            raise ValueError("SummarizeAll.preserve_latest cannot be negative.")
+
+    async def acompress(
+        self,
+        messages: Sequence[BaseMessage],
+        *,
+        runtime: ChatRuntime[TContext],
+        model: BaseChatModel,
+    ) -> list[BaseMessage]:
+        del runtime
+        messages_to_summarize, messages_to_keep = _split_preserved_tail(
+            messages,
+            self.preserve_latest,
+        )
+        if not messages_to_summarize:
+            return _replace_with_latest(messages_to_keep)
+
+        summary = await _summarize_messages(messages_to_summarize, model=model)
+        return _replace_with_summary_between([], summary, messages_to_keep)
+
+
+@dataclass(frozen=True, slots=True)
+class SummarizeLatestN(Generic[TContext]):
+    """Summarize the latest N history messages and preserve the active latest messages."""
+
+    n: int
+    preserve_latest: int = 1
+
+    def __post_init__(self) -> None:
+        if self.n < 1:
+            raise ValueError("SummarizeLatestN.n must be at least 1.")
+        if self.preserve_latest < 0:
+            raise ValueError("SummarizeLatestN.preserve_latest cannot be negative.")
+
+    async def acompress(
+        self,
+        messages: Sequence[BaseMessage],
+        *,
+        runtime: ChatRuntime[TContext],
+        model: BaseChatModel,
+    ) -> list[BaseMessage]:
+        del runtime
+        history_messages, preserved_messages = _split_preserved_tail(
+            messages,
+            self.preserve_latest,
+        )
+        messages_to_summarize = history_messages[-self.n :]
+        messages_to_keep = history_messages[: -self.n]
+        if not messages_to_summarize:
+            return _replace_with_latest(preserved_messages)
+
+        summary = await _summarize_messages(messages_to_summarize, model=model)
+        return _replace_with_summary_between(
+            messages_to_keep,
+            summary,
+            preserved_messages,
+        )

autochat/compression/types.py

@@ -0,0 +1,21 @@
+from collections.abc import Sequence
+from typing import Generic, Protocol, TypeVar
+
+from langchain.chat_models import BaseChatModel
+from langchain_core.messages import BaseMessage
+
+from autochat.runtime import ChatRuntime
+
+TContext = TypeVar("TContext")
+
+
+class CompressionStrategy(Protocol[TContext]):
+    """Strategy used by AutoCompress to compact chat history."""
+
+    async def acompress(
+        self,
+        messages: Sequence[BaseMessage],
+        *,
+        runtime: ChatRuntime[TContext],
+        model: BaseChatModel,
+    ) -> list[BaseMessage]: ...

autochat/config.py

@@ -0,0 +1,11 @@
+from dataclasses import dataclass
+from typing import Any
+
+from langchain.chat_models import BaseChatModel
+
+
+@dataclass(frozen=True, slots=True)
+class ChatConfig:
+    model: BaseChatModel
+    model_kwargs: dict[str, Any] | None = None
+    context_window: int | None = None

autochat/exceptions/__init__.py

@@ -0,0 +1,3 @@
+from .errors import AutoChatError, AutoChatToolError
+
+__all__ = ["AutoChatError", "AutoChatToolError"]

autochat/exceptions/errors.py

@@ -0,0 +1,6 @@
+class AutoChatError(Exception):
+    """Base exception for AutoChat."""
+
+
+class AutoChatToolError(AutoChatError):
+    """Base exception for AutoChat tool execution errors."""

autochat/graph/__init__.py

@@ -0,0 +1,4 @@
+from .builder import build_chat_graph
+from .state import ChatGraphState
+
+__all__ = ["build_chat_graph", "ChatGraphState"]

autochat/graph/builder.py

@@ -0,0 +1,151 @@
+from collections.abc import Sequence
+from typing import Any
+
+from langchain_core.messages import AIMessage, BaseMessage, SystemMessage
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.graph import END, START, StateGraph
+from langgraph.graph.state import CompiledStateGraph
+
+from autochat.compression import AutoCompress
+from autochat.config import ChatConfig
+from autochat.graph.state import ChatGraphState, ChatGraphUpdate
+from autochat.graph.tools import run_tool_calls
+from autochat.graph.runtime import get_runtime
+from autochat.guidelines import ChatGuideline
+from autochat.retrieval import (
+    ChatRetriever,
+    RetrievalConfig,
+)
+from autochat.tools import ChatTool
+
+
+def validate_callable_names(
+    tools: Sequence[ChatTool[Any, Any, Any]],
+    retrievers: Sequence[ChatRetriever[Any]],
+) -> None:
+    """Validate that tool and retriever names are unique."""
+    names = [tool.name for tool in tools] + [retriever.name for retriever in retrievers]
+    duplicates = {name for name in names if names.count(name) > 1}
+
+    if duplicates:
+        duplicate_names = ", ".join(sorted(duplicates))
+        raise ValueError(f"Duplicate tool/retriever names: {duplicate_names}")
+
+
+def build_system_messages(
+    system_message: str | None,
+    guidelines: Sequence[ChatGuideline],
+) -> list[SystemMessage]:
+    messages: list[SystemMessage] = []
+
+    if system_message:
+        messages.append(SystemMessage(content=system_message))
+
+    if guidelines:
+        guideline_text = "\n".join(f"- {guideline.content}" for guideline in guidelines)
+        messages.append(SystemMessage(content=f"Guidelines:\n{guideline_text}"))
+
+    return messages
+
+
+def should_continue(state: ChatGraphState) -> str:
+    last_message = state["messages"][-1]
+
+    if isinstance(last_message, AIMessage) and last_message.tool_calls:
+        return "tools"
+
+    return END
+
+
+def build_chat_graph(
+    *,
+    config: ChatConfig,
+    tools: Sequence[ChatTool[Any, Any, Any]],
+    retrievers: Sequence[ChatRetriever[Any]],
+    retrieval: RetrievalConfig[Any],
+    compression: AutoCompress[Any] | None,
+    system_message: str | None,
+    guidelines: Sequence[ChatGuideline],
+    persistence: BaseCheckpointSaver | None,
+) -> CompiledStateGraph:
+    graph = StateGraph(ChatGraphState)
+    system_messages = build_system_messages(system_message, guidelines)
+
+    chat_config = config
+    model = chat_config.model
+    if tools or retrievers:
+        validate_callable_names(tools, retrievers)
+
+        model = model.bind_tools(
+            [
+                *(tool.model_tool() for tool in tools),
+                *(retriever.model_tool() for retriever in retrievers),
+            ]
+        )
+
+    async def call_model(
+        state: ChatGraphState,
+        config: RunnableConfig,
+    ) -> ChatGraphUpdate:
+        messages: list[BaseMessage] = [
+            *system_messages,
+            *state["messages"],
+        ]
+
+        response = await model.ainvoke(
+            messages, config=config, **(chat_config.model_kwargs or {})
+        )
+
+        response_messages: list[BaseMessage] = [response]
+        return {"messages": response_messages}
+
+    async def call_tools(
+        state: ChatGraphState,
+        config: RunnableConfig,
+    ) -> ChatGraphUpdate:
+        last_message = state["messages"][-1]
+
+        if not isinstance(last_message, AIMessage):
+            return {"messages": []}
+
+        tool_messages = await run_tool_calls(
+            last_message,
+            tools,
+            retrievers,
+            config,
+        )
+        messages: list[BaseMessage] = list(tool_messages)
+        return {"messages": messages}
+
+    async def compress_messages(
+        state: ChatGraphState,
+        config: RunnableConfig,
+    ) -> ChatGraphUpdate:
+        if compression is None:
+            return {"messages": []}
+
+        messages = await compression.acompress_if_needed(
+            list(state["messages"]),
+            runtime=get_runtime(config),
+            config=chat_config,
+        )
+        if messages is None:
+            return {"messages": []}
+
+        return {"messages": messages}
+
+    graph.add_node("agent", call_model)
+    graph.add_node("tools", call_tools)
+
+    if compression:
+        graph.add_node("compression", compress_messages)
+        graph.add_edge(START, "compression")
+        graph.add_edge("compression", "agent")
+    else:
+        graph.add_edge(START, "agent")
+
+    graph.add_conditional_edges("agent", should_continue)
+    graph.add_edge("tools", "agent")
+
+    return graph.compile(checkpointer=persistence)