Add ConversationStore (#96)

Author: mateusz834

splunklib/ai/README.md

@@ -264,6 +264,83 @@ async with Agent(
 ) as agent: ...
 ```
 
+## Conversation stores
+
+By default, each call to `agent.invoke` is stateless - the agent has no memory of previous interactions,
+unless you provide the previouis message history explicitly. A conversation store enables the agent to persist
+and recall message history across invocations.
+
+### `InMemoryStore`
+
+The built-in `InMemoryStore` keeps conversation history in process memory.
+
+```py
+from splunklib.ai import Agent, OpenAIModel
+from splunklib.ai.conversation_store import InMemoryStore
+from splunklib.ai.messages import HumanMessage
+from splunklib.client import connect
+
+model = OpenAIModel(...)
+service = connect(...)
+
+async with Agent(
+    model=model,
+    service=service,
+    system_prompt="",
+    conversation_store=InMemoryStore(),
+) as agent:
+    await agent.invoke([HumanMessage(content="Hi, my name is Chris.")])
+    result = await agent.invoke([HumanMessage(content="What is my name?")])
+    print(result.final_message.content)  # Chris
+```
+
+### Multiple conversation threads
+
+Each conversation is isolated by a `thread_id`. You can pass a `thread_id` per invocation to maintain
+separate histories for different users or sessions within the same agent instance.
+
+```py
+async with Agent(
+    model=model,
+    service=service,
+    system_prompt="",
+    conversation_store=InMemoryStore(),
+) as agent:
+    await agent.invoke(
+        [HumanMessage(content="Hi, my name is Alice.")],
+        thread_id="user-alice",
+    )
+    await agent.invoke(
+        [HumanMessage(content="Hi, my name is Bob.")],
+        thread_id="user-bob",
+    )
+
+    result = await agent.invoke(
+        [HumanMessage(content="What is my name?")],
+        thread_id="user-alice",
+    )
+    print(result.final_message.content)  # Alice - Bob's thread is unaffected
+```
+
+A custom `thread_id` can also be set on the agent constructor. When `invoke` is called without an explicit
+`thread_id`, the `thread_id` from the constructor is used. If no `thread_id` is provided in the constructor, one
+is generated implicitly.
+
+```py
+async with Agent(
+    model=model,
+    service=service,
+    system_prompt="",
+    conversation_store=InMemoryStore(),
+    thread_id="session-42",
+) as agent:
+    await agent.invoke([HumanMessage(content="Hi, my name is Chris.")])
+
+    # No thread_id supplied — falls back to "session-42"
+    result = await agent.invoke([HumanMessage(content="What is my name?")])
+    print(result.final_message.content)  # Chris
+```
+
 ## Subagents
 
 The `Agent` constructor can accept subagents as input parameters.

splunklib/ai/agent.py

@@ -17,10 +17,12 @@
 from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
 from logging import Logger
 from typing import Self, final, override
+from uuid import uuid4
 
 from pydantic import BaseModel
 
 from splunklib.ai.base_agent import BaseAgent
+from splunklib.ai.conversation_store import ConversationStore
 from splunklib.ai.core.backend import AgentImpl
 from splunklib.ai.core.backend_registry import get_backend
 from splunklib.ai.messages import AgentResponse, BaseMessage, OutputT
@@ -107,6 +109,29 @@ class Agent(BaseAgent[OutputT]):
         logger:
             Optional logger instance used for tracing and debugging the agent's execution.
             Additionally logs from the local tools are forwarded to this logger.
+
+        conversation_store:
+            Optional `ConversationStore` instance used to persist conversation history
+            across multiple `invoke` calls. When provided, the agent automatically loads
+            prior messages for the active thread before each invocation and saves the
+            full updated history afterwards.
+
+            Use the built-in `InMemoryStore` for in-process persistence, or implement
+            `ConversationStore` to back history with an external store.
+
+            Without a store, each `invoke` call is stateless and the agent has no memory
+            of previous turns.
+
+        thread_id:
+            Identifies the conversation thread used when reading from and writing to the
+            `conversation_store`. Each unique `thread_id` maintains a separate history,
+            so different users or sessions can share one store without interference.
+
+            If omitted, a random ID is generated automatically. The `thread_id` can
+            also be overridden per-call by passing it directly to `invoke`.
+
+            Never invoke an Agent using the same thread_id more than once concurrently
+            while using the same conversation_store.
     """
 
     _impl: AgentImpl[OutputT] | None
@@ -129,17 +154,22 @@ def __init__(
         name: str = "",  # Only used by Subagents
         description: str = "",  # Only used by Subagents
         logger: Logger | None = None,
+        conversation_store: ConversationStore | None = None,
+        thread_id: str | None = None,
     ) -> None:
         super().__init__(
             model=model,
             system_prompt=system_prompt,
             name=name,
             description=description,
+            tools=None,
             agents=agents,
             input_schema=input_schema,
             output_schema=output_schema,
             middleware=middleware,
             logger=logger,
+            conversation_store=conversation_store,
+            thread_id=thread_id if thread_id is not None else str(uuid4()),
         )
 
         self._use_mcp_tools = use_mcp_tools
@@ -242,12 +272,19 @@ async def __aexit__(
         self._agent_context_manager = None
         return result
 
+    # TODO: for now we have a thread_id as an optional param, should
+    # we wrap it in a dataclass? Might help with future-proofing the API??
     @override
-    async def invoke(self, messages: list[BaseMessage]) -> AgentResponse[OutputT]:
+    async def invoke(
+        self, messages: list[BaseMessage], thread_id: str | None = None
+    ) -> AgentResponse[OutputT]:
         if not self._impl:
             raise AssertionError("Agent must be used inside 'async with'")
 
-        return await self._impl.invoke(messages)
+        if thread_id is None:
+            thread_id = self._thread_id
+
+        return await self._impl.invoke(messages, thread_id)
 
 
 def _local_tools_path() -> tuple[str | None, str]:

splunklib/ai/base_agent.py

@@ -20,6 +20,7 @@
 
 from pydantic import BaseModel
 
+from splunklib.ai.conversation_store import ConversationStore
 from splunklib.ai.messages import AgentResponse, BaseMessage, OutputT
 from splunklib.ai.middleware import AgentMiddleware
 from splunklib.ai.model import PredefinedModel
@@ -38,19 +39,23 @@ class BaseAgent(Generic[OutputT], ABC):
     _middleware: Sequence[AgentMiddleware] | None = None
     _trace_id: str
     _logger: logging.Logger
+    _conversation_store: ConversationStore | None = None
+    _thread_id: str
 
     def __init__(
         self,
         system_prompt: str,
         model: PredefinedModel,
-        description: str = "",
-        name: str = "",
-        tools: Sequence[Tool] | None = None,
-        agents: Sequence["BaseAgent[BaseModel | None]"] | None = None,
-        input_schema: type[BaseModel] | None = None,
-        output_schema: type[OutputT] | None = None,
-        middleware: Sequence[AgentMiddleware] | None = None,
-        logger: logging.Logger | None = None,
+        description: str,
+        name: str,
+        tools: Sequence[Tool] | None,
+        agents: Sequence["BaseAgent[BaseModel | None]"] | None,
+        input_schema: type[BaseModel] | None,
+        output_schema: type[OutputT] | None,
+        middleware: Sequence[AgentMiddleware] | None,
+        logger: logging.Logger | None,
+        conversation_store: ConversationStore | None,
+        thread_id: str,
     ) -> None:
         self._system_prompt = system_prompt
         self._model = model
@@ -62,6 +67,8 @@ def __init__(
         self._output_schema = output_schema
         self._middleware = tuple(middleware) if middleware else ()
         self._trace_id = secrets.token_hex(16)  # 32 Hex characters
+        self._conversation_store = conversation_store
+        self._thread_id = thread_id
 
         if logger is None:
             # Create a no-op logger to skip checking for its existence.
@@ -70,7 +77,9 @@ def __init__(
         self._logger = logger
 
     @abstractmethod
-    async def invoke(self, messages: list[BaseMessage]) -> AgentResponse[OutputT]: ...
+    async def invoke(
+        self, messages: list[BaseMessage], thread_id: str | None = None
+    ) -> AgentResponse[OutputT]: ...
 
     @property
     def logger(self) -> logging.Logger:
@@ -115,3 +124,11 @@ def middleware(self) -> Sequence[AgentMiddleware] | None:
     @property
     def trace_id(self) -> str:
         return self._trace_id
+
+    @property
+    def conversation_store(self) -> ConversationStore | None:
+        return self._conversation_store
+
+    @property
+    def default_thread_id(self) -> str:
+        return self._thread_id

splunklib/ai/conversation_store.py

@@ -0,0 +1,41 @@
+# Copyright © 2011-2026 Splunk, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"): you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+from collections.abc import Sequence
+from typing import Protocol, override
+
+from splunklib.ai.messages import BaseMessage
+
+
+class ConversationStore(Protocol):
+    async def get_messages(self, thread_id: str) -> Sequence[BaseMessage]: ...
+
+    async def store_messages(
+        self, thread_id: str, messages: list[BaseMessage]
+    ) -> None: ...
+
+
+class InMemoryStore(ConversationStore):
+    _threads: dict[str, Sequence[BaseMessage]]
+
+    def __init__(self) -> None:
+        self._threads = {}
+
+    @override
+    async def get_messages(self, thread_id: str) -> Sequence[BaseMessage]:
+        return self._threads.get(thread_id, [])
+
+    @override
+    async def store_messages(self, thread_id: str, messages: list[BaseMessage]) -> None:
+        self._threads[thread_id] = messages.copy()

splunklib/ai/core/backend.py

@@ -29,7 +29,9 @@ class InvalidMessageTypeError(Exception):
 class AgentImpl(Protocol[OutputT]):
     """Backend-specific agent implementation used by the public `Agent` wrapper."""
 
-    async def invoke(self, messages: list[BaseMessage]) -> AgentResponse[OutputT]: ...
+    async def invoke(
+        self, messages: list[BaseMessage], thread_id: str
+    ) -> AgentResponse[OutputT]: ...
 
 
 class Backend(Protocol):