soliplex 0.42__py3-none-any.whl

soliplex/agents.py

@@ -0,0 +1,137 @@
+import dataclasses
+import typing
+from collections import abc
+
+import pydantic_ai
+from pydantic_ai import agent as ai_agent
+from pydantic_ai import mcp as ai_mcp
+from pydantic_ai import tools as ai_tools
+from pydantic_ai.models import google as google_models
+from pydantic_ai.models import openai as openai_models
+from pydantic_ai.providers import google as google_providers
+from pydantic_ai.providers import ollama as ollama_providers
+from pydantic_ai.providers import openai as openai_providers
+
+from soliplex import agui
+from soliplex import config
+from soliplex import mcp_client
+from soliplex import models
+
+ToolConfigMap = dict[str, typing.Any]
+
+
+@dataclasses.dataclass
+class AgentDependencies:
+    """Agent dependencies implementing StateHandler protocol.
+
+    The `state` field is required by pydantic-ai's StateHandler protocol.
+    AG-UI will inject the client's state into this field for each run.
+    """
+
+    the_installation: typing.Any  # installation.Installation
+    user: models.UserProfile = None  # TBD make required
+    tool_configs: ToolConfigMap = None
+    state: agui.AGUI_State = dataclasses.field(default_factory=dict)
+
+
+SoliplexAgent = ai_agent.AbstractAgent[AgentDependencies, typing.Any]
+AgentFactory = abc.Callable[
+    [
+        config.AgentConfig,
+        ToolConfigMap,
+        config.MCP_ClientToolsetConfigMap,
+    ],
+    SoliplexAgent,
+]
+
+# Cache for agents to avoid recreating them
+_agent_cache: dict[str, pydantic_ai.Agent] = {}
+
+
+def make_ai_tool(tool_config: config.ToolConfig) -> ai_tools.Tool:
+    tool_func = tool_config.tool_with_config
+
+    return ai_tools.Tool(
+        tool_func,
+        name=tool_config.tool_id,
+    )
+
+
+def make_mcp_client_toolset(
+    toolset_config: config.MCP_ClientToolsetConfig,
+) -> ai_mcp.MCPServer:
+    toolset_klass = mcp_client.TOOLSET_CLASS_BY_KIND[toolset_config.kind]
+    return toolset_klass(**toolset_config.tool_kwargs)
+
+
+def _get_default_agent_from_configs(
+    agent_config: config.AgentConfig,
+    tool_configs: ToolConfigMap,
+    mcp_client_toolset_configs: config.MCP_ClientToolsetConfigMap,
+) -> SoliplexAgent:
+    """Build a Pydantic AI agent from a config"""
+    provider_kw = agent_config.llm_provider_kw
+
+    if agent_config.provider_type == config.LLMProviderType.GOOGLE:
+        provider = google_providers.GoogleProvider(**provider_kw)
+        model = google_models.GoogleModel(
+            model_name=agent_config.model_name,
+            provider=provider,
+        )
+
+    elif agent_config.provider_type == config.LLMProviderType.OLLAMA:
+        provider_kw["api_key"] = "dummy"
+        provider = ollama_providers.OllamaProvider(**provider_kw)
+        model = openai_models.OpenAIChatModel(
+            model_name=agent_config.model_name,
+            provider=provider,
+        )
+    else:
+        provider = openai_providers.OpenAIProvider(**provider_kw)
+        model = openai_models.OpenAIChatModel(
+            model_name=agent_config.model_name,
+            provider=provider,
+        )
+
+    tools = [
+        make_ai_tool(tool_config) for tool_config in tool_configs.values()
+    ]
+    toolsets = [
+        make_mcp_client_toolset(mctc)
+        for mctc in mcp_client_toolset_configs.values()
+    ]
+
+    return pydantic_ai.Agent(
+        model=model,
+        model_settings=agent_config.model_settings,
+        tools=tools,
+        toolsets=toolsets,
+        instructions=agent_config.get_system_prompt(),
+        deps_type=AgentDependencies,
+    )
+
+
+def get_agent_from_configs(
+    agent_config: config.AgentConfig,
+    tool_configs: ToolConfigMap,
+    mcp_client_toolset_configs: config.MCP_ClientToolsetConfigMap,
+) -> SoliplexAgent:
+    """Get or create an agent from the specified agent and tool configs."""
+
+    if agent_config.id not in _agent_cache:
+        if agent_config.kind == "default":
+            agent = _get_default_agent_from_configs(
+                agent_config,
+                tool_configs,
+                mcp_client_toolset_configs,
+            )
+
+        else:
+            agent = agent_config.factory(
+                tool_configs=tool_configs,
+                mcp_client_toolset_configs=mcp_client_toolset_configs,
+            )
+
+        _agent_cache[agent_config.id] = agent
+
+    return _agent_cache[agent_config.id]

soliplex/agui/__init__.py

@@ -0,0 +1,380 @@
+from __future__ import annotations
+
+import abc
+import collections.abc
+import datetime
+import typing
+
+import fastapi
+from ag_ui import core as agui_core
+from sqlalchemy.ext import asyncio as sqla_asyncio
+
+AGUI_Events = list[agui_core.Event]
+AGUI_EventStream = collections.abc.AsyncIterator[agui_core.Event]
+AGUI_State = dict[str, typing.Any]
+
+
+_COMPACTIBLE_TYPES = {
+    agui_core.EventType.TEXT_MESSAGE_CONTENT,
+    agui_core.EventType.THINKING_TEXT_MESSAGE_CONTENT,
+}
+
+
+async def compact_event_stream(stream: AGUI_EventStream):
+    compacting: agui_core.Event = None
+    compacting_id: str = None
+
+    async for event in stream:
+        if compacting is not None:
+            event_id = getattr(event, "message_id", None)
+            if event.type == compacting.type and event_id == compacting_id:
+                compacting.delta += event.delta
+            else:
+                to_yield, compacting = compacting, None
+                yield to_yield
+                yield event
+
+        else:
+            if event.type in _COMPACTIBLE_TYPES:
+                compacting = event.model_copy()
+                compacting_id = getattr(event, "message_id", None)
+            else:
+                yield event
+
+
+class AGUI_Exception(ValueError):
+    status_code = 400
+
+
+class UnknownThread(AGUI_Exception):
+    status_code = 404
+
+    def __init__(self, user_name: str, thread_id: str):
+        self.user_name = user_name
+        self.thread_id = thread_id
+        message = f"Unknown thread: UUID {thread_id} for user {user_name}"
+        super().__init__(message)
+
+
+class UnknownRun(AGUI_Exception):
+    status_code = 404
+
+    def __init__(self, run_id: str):
+        self.run_id = run_id
+        super().__init__(
+            f"Unknown run: UUID {run_id} does not exist in thread"
+        )
+
+
+class ThreadRoomMismatch(AGUI_Exception):
+    def __init__(self, room_id: str, thread_room_id: str):
+        self.room_id = room_id
+        self.thread_room_id = thread_room_id
+        super().__init__(
+            f"Thread room ID '{thread_room_id}' "
+            f"does not match room ID '{room_id}'"
+        )
+
+
+class MissingParentRun(AGUI_Exception):
+    def __init__(self, parent_run_id: str):
+        self.parent_run_id = parent_run_id
+        super().__init__(
+            f"Unknown parent run: UUID {parent_run_id} "
+            f"does not exist in thread"
+        )
+
+
+class RunAlreadyStarted(AGUI_Exception):
+    def __init__(self, user_name: str, thread_id: str, run_id: str):
+        self.user_name = user_name
+        self.thread_id = thread_id
+        self.run_id = run_id
+        super().__init__(f"Run already started: UUID {run_id}")
+
+
+#
+#   ABCs defined here are notional contracts.
+#
+
+
+RunUsageStats = collections.namedtuple(
+    "RunUsageStats",
+    [
+        "input_tokens",
+        "output_tokens",
+        "requests",
+        "tool_calls",
+    ],
+)
+
+
+class RunUsage(abc.ABC):
+    """LLM usage for a run"""
+
+    input_tokens: int
+    """LLM input tokens consumed"""
+
+    output_tokens: int
+    """LLM output tokens consumed"""
+
+    requests: int
+    """LLM requests made"""
+
+    tool_calls: int
+    """LLM tool_calls made"""
+
+    @abc.abstractmethod
+    def as_tuple(self) -> RunUsageStats:
+        """Return values as a tuple."""
+
+
+class RunMetadata(abc.ABC):
+    """User-defined metadata for a run"""
+
+    label: str
+    """Label for a run (similar to a git tag)"""
+
+
+class Run(abc.ABC):
+    """Input data and events for an AGUI run
+
+    Runs are not accessed directly:  use 'ThreadStorage'.
+    """
+
+    thread_id: str
+    """ID of the thread in which the run was created"""
+
+    run_id: str
+    """Unique ID for a run"""
+
+    parent_run_id: str | None
+    """ID of the parent run"""
+
+    run_usage: RunUsage | None
+    """Optional LLM usage data for a run"""
+
+    run_metadata: RunMetadata | None
+    """Optional user-defined metadata for a run"""
+
+    run_input: agui_core.RunAgentInput
+    """Input from the client-request which initiates the AG-UI run"""
+
+    created: datetime.datetime
+    """Timestamp"""
+
+    @abc.abstractmethod
+    async def list_events(self) -> AGUI_Events:
+        """Return AGUI events for the run"""
+
+
+class ThreadMetadata(abc.ABC):
+    """Optional user-defined thread metadata"""
+
+    name: str
+    """Name for the thread"""
+
+    description: str
+    """Description for the thread"""
+
+
+class Thread(abc.ABC):
+    """Hold a set of AGUI runs sharing the same 'thread_id'
+
+    Runs are not accessed directly:  use 'ThreadStorage'.
+    """
+
+    thread_id: str
+    """Unique ID for the thread"""
+
+    room_id: str
+    """ID for room in which the thread was created"""
+
+    thread_metadata: ThreadMetadata | None
+    """Optional thread metadata"""
+
+    created: datetime.datetime
+    """Timestamp"""
+
+    @abc.abstractmethod
+    async def list_runs(self) -> list[Run]:
+        """Return runs for this thread"""
+
+
+class ThreadStorage(abc.ABC):
+    @abc.abstractmethod
+    async def list_user_threads(
+        self,
+        *,
+        user_name: str,
+        room_id: str = None,
+    ) -> list[Thread]:
+        """Return a list of the user's threads.
+
+        If 'room_id' is passed, filter the threads to those created
+        in that room.
+        """
+
+    @abc.abstractmethod
+    async def get_thread(
+        self,
+        *,
+        user_name: str,
+        room_id: str,
+        thread_id: str,
+    ) -> Thread:
+        """Return the actual thread instance
+
+        N.B.:  caller must treat the instance as read-only!
+        """
+
+    @abc.abstractmethod
+    async def new_thread(
+        self,
+        *,
+        user_name: str,
+        room_id: str,
+        thread_metadata: ThreadMetadata | dict = None,
+        initial_run: bool = True,
+    ) -> Thread:
+        """Create a new thread
+
+        If 'thread_metadata' is a dict, convert it to a 'ThreadMetadata'
+        instance.
+        """
+
+    @abc.abstractmethod
+    async def update_thread_metadata(
+        self,
+        *,
+        user_name: str,
+        room_id: str,
+        thread_id: str,
+        thread_metadata: ThreadMetadata | dict = None,
+    ) -> Thread:
+        """Update thread instance with the given metadata, or None
+
+        If 'thread_metadata' is a dict, convert it to a 'ThreadMetadata'
+        instance.
+
+        If 'thread_metadata' is None, or an empty dict, remove any existing
+        metadata on the thread.
+        """
+
+    @abc.abstractmethod
+    async def delete_thread(
+        self,
+        *,
+        user_name: str,
+        room_id: str,
+        thread_id: str,
+    ) -> None:
+        """Remove a thread"""
+
+    @abc.abstractmethod
+    async def new_run(
+        self,
+        *,
+        user_name: str,
+        room_id: str,
+        thread_id: str,
+        run_metadata: RunMetadata = None,
+        parent_run_id: str = None,
+    ) -> Run:
+        """Create a new run for the thread
+
+        If 'run_metadata' is a dict, convert it to a 'RunMetadata' instance.
+
+        If 'parent_run_id' is passed, ensure it is valid.
+        """
+
+    @abc.abstractmethod
+    async def get_run(
+        self,
+        user_name: str,
+        room_id: str,
+        thread_id: str,
+        run_id: str,
+    ) -> Run:
+        """Return an existing run for a thread"""
+
+    @abc.abstractmethod
+    async def add_run_input(
+        self,
+        *,
+        user_name: str,
+        room_id: str,
+        thread_id: str,
+        run_id: str,
+        run_input: agui_core.RunAgentInput,
+    ) -> Run:
+        """Update a run with the given 'run_agent_input'"""
+
+    @abc.abstractmethod
+    async def update_run_metadata(
+        self,
+        *,
+        user_name: str,
+        room_id: str,
+        thread_id: str,
+        run_id: str,
+        run_metadata: RunMetadata | dict = None,
+    ) -> Run:
+        """Update a run with the given metadata
+
+        If 'run_metadata' is a dict, convert it to a 'RunMetadata' instance.
+
+        If 'run_metadata' is None, or an empty dict, remove any existing
+        metadata on the run.
+        """
+
+    @abc.abstractmethod
+    async def save_run_events(
+        self,
+        *,
+        user_name: str,
+        room_id: str,
+        thread_id: str,
+        run_id: str,
+        events: AGUI_Events,
+    ) -> AGUI_Events:
+        """Save the events for a gven run"""
+
+    @abc.abstractmethod
+    async def save_run_usage(
+        self,
+        *,
+        user_name: str,
+        room_id: str,
+        thread_id: str,
+        run_id: str,
+        input_tokens: int,
+        output_tokens: int,
+        requests: int,
+        tool_calls: int,
+    ):
+        """Save the run usage statistics"""
+
+    @abc.abstractmethod
+    async def save_run_feedback(
+        self,
+        *,
+        user_name: str,
+        room_id: str,
+        thread_id: str,
+        run_id: str,
+        feedback: str,
+        reason: str,
+    ):
+        """Save the run feedback"""
+
+
+async def get_the_threads(request: fastapi.Request) -> ThreadStorage:
+    from . import persistence
+
+    engine = request.state.threads_engine
+    async with sqla_asyncio.AsyncSession(bind=engine) as session:
+        yield persistence.ThreadStorage(session)
+
+
+depend_the_threads = fastapi.Depends(get_the_threads)

soliplex/agui/features.py

@@ -0,0 +1,47 @@
+"""Features defined by soliplex"""
+
+import pydantic
+from haiku.rag.agents.chat import state as hr_chat_state
+from haiku.rag.agents.research import models as hr_graph_models
+
+FILTER_DOCUMENTS_FEATURE = "filter_documents"
+ASK_HISTORY_FEATURE = "ask_history"
+HAIKU_CHAT_FEATURE = hr_chat_state.AGUI_STATE_KEY
+
+KW_ONLY_NONE = pydantic.Field(kw_only=True, default=None)
+
+
+class FilterDocuments(pydantic.BaseModel):
+    """Documents selected by the user to be used to answer a question
+
+    This model describes the 'filter_documents' key in the AG-UI state.
+
+    If 'document_ids' is empty or None, or if the 'filter_documents'
+    key is not present in the AG-UI state, no filter is applied:  the
+    'ask_with_rich_citations' tool will return all documents matching
+    the query from the LLM.
+    """
+
+    document_ids: list[str] | None = KW_ONLY_NONE
+
+
+class QuestionResponseCitations(pydantic.BaseModel):
+    """Single question to the 'ask_with_rich_citations' tool
+
+    Includes response generated by the tool and accompanying citations.
+    """
+
+    question: str
+    response: str
+    citations: list[hr_graph_models.Citation] = []
+
+
+class AskedAndAnswered(pydantic.BaseModel):
+    """History of questions asked and their replies + citations
+
+    This model describes the 'ask_history' key in the AG-UI state.
+
+    Each call to the 'ask_with_rich_citations' took append
+    """
+
+    questions: list[QuestionResponseCitations] = []