plain-agent 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (35) hide show
  1. plain_agent-0.1.0/LICENSE +21 -0
  2. plain_agent-0.1.0/PKG-INFO +64 -0
  3. plain_agent-0.1.0/README.md +52 -0
  4. plain_agent-0.1.0/plain_agent.egg-info/PKG-INFO +64 -0
  5. plain_agent-0.1.0/plain_agent.egg-info/SOURCES.txt +33 -0
  6. plain_agent-0.1.0/plain_agent.egg-info/dependency_links.txt +1 -0
  7. plain_agent-0.1.0/plain_agent.egg-info/entry_points.txt +2 -0
  8. plain_agent-0.1.0/plain_agent.egg-info/requires.txt +2 -0
  9. plain_agent-0.1.0/plain_agent.egg-info/top_level.txt +1 -0
  10. plain_agent-0.1.0/pyproject.toml +22 -0
  11. plain_agent-0.1.0/setup.cfg +4 -0
  12. plain_agent-0.1.0/simple_agent/agent_loop.py +112 -0
  13. plain_agent-0.1.0/simple_agent/cli.py +43 -0
  14. plain_agent-0.1.0/simple_agent/conversation_history.py +59 -0
  15. plain_agent-0.1.0/simple_agent/llm_client.py +49 -0
  16. plain_agent-0.1.0/simple_agent/message_types.py +39 -0
  17. plain_agent-0.1.0/simple_agent/prompt.py +5 -0
  18. plain_agent-0.1.0/simple_agent/streaming.py +96 -0
  19. plain_agent-0.1.0/simple_agent/terminal_loop.py +49 -0
  20. plain_agent-0.1.0/simple_agent/tools/base_tool.py +29 -0
  21. plain_agent-0.1.0/simple_agent/tools/command_policy.py +20 -0
  22. plain_agent-0.1.0/simple_agent/tools/command_runtime.py +155 -0
  23. plain_agent-0.1.0/simple_agent/tools/edit_file.py +75 -0
  24. plain_agent-0.1.0/simple_agent/tools/list_files.py +45 -0
  25. plain_agent-0.1.0/simple_agent/tools/permissions/file_permission.py +94 -0
  26. plain_agent-0.1.0/simple_agent/tools/read_file.py +60 -0
  27. plain_agent-0.1.0/simple_agent/tools/run_command.py +54 -0
  28. plain_agent-0.1.0/simple_agent/tools/search_text.py +66 -0
  29. plain_agent-0.1.0/simple_agent/tools/tools.py +46 -0
  30. plain_agent-0.1.0/simple_agent/tools/utils.py +11 -0
  31. plain_agent-0.1.0/simple_agent/tools/write_file.py +55 -0
  32. plain_agent-0.1.0/tests/test_agent.py +376 -0
  33. plain_agent-0.1.0/tests/test_llm_client.py +78 -0
  34. plain_agent-0.1.0/tests/test_terminal_loop.py +68 -0
  35. plain_agent-0.1.0/tests/test_tools.py +369 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Ribo Mo
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,64 @@
1
+ Metadata-Version: 2.4
2
+ Name: plain-agent
3
+ Version: 0.1.0
4
+ Summary: A small agent that calls OpenAI-compatible APIs.
5
+ License-Expression: MIT
6
+ Requires-Python: >=3.11
7
+ Description-Content-Type: text/markdown
8
+ License-File: LICENSE
9
+ Requires-Dist: openai==2.43.0
10
+ Requires-Dist: python-dotenv>=1.2.2
11
+ Dynamic: license-file
12
+
13
+ # Plain Agent
14
+
15
+ A small agent that calls OpenAI compatible LLM APIs.
16
+
17
+ This project starts with a streaming agent loop:
18
+
19
+ ```text
20
+ Interactive terminal
21
+ -> reads your prompt
22
+ -> streams assistant text as it arrives
23
+ -> detects tool calls from the model
24
+ -> runs workspace tools when requested
25
+ -> sends tool results back to the model
26
+ -> repeats until the assistant gives a final answer
27
+ ```
28
+
29
+ ## Install
30
+
31
+ This project uses `uv` to track the Python environment.
32
+ If `uv` is not installed, follow the [official installation guide](https://docs.astral.sh/uv/getting-started/installation/).
33
+
34
+ ```bash
35
+ uv sync
36
+ ```
37
+
38
+ ## Configuration
39
+
40
+ Create a local `.env` file or export environment variables in your shell. See `.env.example` for more examples.
41
+
42
+ For DeepSeek:
43
+
44
+ ```bash
45
+ export DEEPSEEK_API_KEY="your-api-key"
46
+ export LLM_PROVIDER="deepseek"
47
+ export LLM_MODEL="deepseek-v4-flash"
48
+ ```
49
+
50
+ For OpenAI:
51
+
52
+ ```bash
53
+ export OPENAI_API_KEY="your-api-key"
54
+ export LLM_PROVIDER="openai"
55
+ export LLM_MODEL="gpt-5.4-mini"
56
+ ```
57
+
58
+ You can still set `LLM_BASE_URL` when you want to override the provider default, such as pointing at a local OpenAI compatible server like Ollama.
59
+
60
+ ## Run
61
+
62
+ ```bash
63
+ uv run plain-agent
64
+ ```
@@ -0,0 +1,52 @@
1
+ # Plain Agent
2
+
3
+ A small agent that calls OpenAI compatible LLM APIs.
4
+
5
+ This project starts with a streaming agent loop:
6
+
7
+ ```text
8
+ Interactive terminal
9
+ -> reads your prompt
10
+ -> streams assistant text as it arrives
11
+ -> detects tool calls from the model
12
+ -> runs workspace tools when requested
13
+ -> sends tool results back to the model
14
+ -> repeats until the assistant gives a final answer
15
+ ```
16
+
17
+ ## Install
18
+
19
+ This project uses `uv` to track the Python environment.
20
+ If `uv` is not installed, follow the [official installation guide](https://docs.astral.sh/uv/getting-started/installation/).
21
+
22
+ ```bash
23
+ uv sync
24
+ ```
25
+
26
+ ## Configuration
27
+
28
+ Create a local `.env` file or export environment variables in your shell. See `.env.example` for more examples.
29
+
30
+ For DeepSeek:
31
+
32
+ ```bash
33
+ export DEEPSEEK_API_KEY="your-api-key"
34
+ export LLM_PROVIDER="deepseek"
35
+ export LLM_MODEL="deepseek-v4-flash"
36
+ ```
37
+
38
+ For OpenAI:
39
+
40
+ ```bash
41
+ export OPENAI_API_KEY="your-api-key"
42
+ export LLM_PROVIDER="openai"
43
+ export LLM_MODEL="gpt-5.4-mini"
44
+ ```
45
+
46
+ You can still set `LLM_BASE_URL` when you want to override the provider default, such as pointing at a local OpenAI compatible server like Ollama.
47
+
48
+ ## Run
49
+
50
+ ```bash
51
+ uv run plain-agent
52
+ ```
@@ -0,0 +1,64 @@
1
+ Metadata-Version: 2.4
2
+ Name: plain-agent
3
+ Version: 0.1.0
4
+ Summary: A small agent that calls OpenAI-compatible APIs.
5
+ License-Expression: MIT
6
+ Requires-Python: >=3.11
7
+ Description-Content-Type: text/markdown
8
+ License-File: LICENSE
9
+ Requires-Dist: openai==2.43.0
10
+ Requires-Dist: python-dotenv>=1.2.2
11
+ Dynamic: license-file
12
+
13
+ # Plain Agent
14
+
15
+ A small agent that calls OpenAI compatible LLM APIs.
16
+
17
+ This project starts with a streaming agent loop:
18
+
19
+ ```text
20
+ Interactive terminal
21
+ -> reads your prompt
22
+ -> streams assistant text as it arrives
23
+ -> detects tool calls from the model
24
+ -> runs workspace tools when requested
25
+ -> sends tool results back to the model
26
+ -> repeats until the assistant gives a final answer
27
+ ```
28
+
29
+ ## Install
30
+
31
+ This project uses `uv` to track the Python environment.
32
+ If `uv` is not installed, follow the [official installation guide](https://docs.astral.sh/uv/getting-started/installation/).
33
+
34
+ ```bash
35
+ uv sync
36
+ ```
37
+
38
+ ## Configuration
39
+
40
+ Create a local `.env` file or export environment variables in your shell. See `.env.example` for more examples.
41
+
42
+ For DeepSeek:
43
+
44
+ ```bash
45
+ export DEEPSEEK_API_KEY="your-api-key"
46
+ export LLM_PROVIDER="deepseek"
47
+ export LLM_MODEL="deepseek-v4-flash"
48
+ ```
49
+
50
+ For OpenAI:
51
+
52
+ ```bash
53
+ export OPENAI_API_KEY="your-api-key"
54
+ export LLM_PROVIDER="openai"
55
+ export LLM_MODEL="gpt-5.4-mini"
56
+ ```
57
+
58
+ You can still set `LLM_BASE_URL` when you want to override the provider default, such as pointing at a local OpenAI compatible server like Ollama.
59
+
60
+ ## Run
61
+
62
+ ```bash
63
+ uv run plain-agent
64
+ ```
@@ -0,0 +1,33 @@
1
+ LICENSE
2
+ README.md
3
+ pyproject.toml
4
+ plain_agent.egg-info/PKG-INFO
5
+ plain_agent.egg-info/SOURCES.txt
6
+ plain_agent.egg-info/dependency_links.txt
7
+ plain_agent.egg-info/entry_points.txt
8
+ plain_agent.egg-info/requires.txt
9
+ plain_agent.egg-info/top_level.txt
10
+ simple_agent/agent_loop.py
11
+ simple_agent/cli.py
12
+ simple_agent/conversation_history.py
13
+ simple_agent/llm_client.py
14
+ simple_agent/message_types.py
15
+ simple_agent/prompt.py
16
+ simple_agent/streaming.py
17
+ simple_agent/terminal_loop.py
18
+ simple_agent/tools/base_tool.py
19
+ simple_agent/tools/command_policy.py
20
+ simple_agent/tools/command_runtime.py
21
+ simple_agent/tools/edit_file.py
22
+ simple_agent/tools/list_files.py
23
+ simple_agent/tools/read_file.py
24
+ simple_agent/tools/run_command.py
25
+ simple_agent/tools/search_text.py
26
+ simple_agent/tools/tools.py
27
+ simple_agent/tools/utils.py
28
+ simple_agent/tools/write_file.py
29
+ simple_agent/tools/permissions/file_permission.py
30
+ tests/test_agent.py
31
+ tests/test_llm_client.py
32
+ tests/test_terminal_loop.py
33
+ tests/test_tools.py
@@ -0,0 +1,2 @@
1
+ [console_scripts]
2
+ plain-agent = simple_agent.cli:main
@@ -0,0 +1,2 @@
1
+ openai==2.43.0
2
+ python-dotenv>=1.2.2
@@ -0,0 +1 @@
1
+ simple_agent
@@ -0,0 +1,22 @@
1
+ [build-system]
2
+ requires = ["setuptools>=77.0.3"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "plain-agent"
7
+ version = "0.1.0"
8
+ description = "A small agent that calls OpenAI-compatible APIs."
9
+ readme = "README.md"
10
+ requires-python = ">=3.11"
11
+ license = "MIT"
12
+ license-files = ["LICENSE"]
13
+ dependencies = [
14
+ "openai==2.43.0",
15
+ "python-dotenv>=1.2.2",
16
+ ]
17
+
18
+ [project.scripts]
19
+ plain-agent = "simple_agent.cli:main"
20
+
21
+ [tool.setuptools.packages.find]
22
+ include = ["simple_agent*"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,112 @@
1
+ """Simple tool loop for the educational agent."""
2
+
3
+ from collections.abc import Iterable
4
+ import json
5
+ from typing import Any, Callable, Generator
6
+
7
+ from simple_agent.conversation_history import ConversationHistory
8
+ from simple_agent.message_types import ToolCallDict
9
+ from simple_agent.prompt import INITIAL_PROMPT
10
+ from simple_agent.streaming import (
11
+ ChatCompletionStreamAccumulator,
12
+ TextDelta,
13
+ ToolResult,
14
+ )
15
+ from simple_agent.tools.tools import Tools
16
+ from simple_agent.tools.utils import error
17
+
18
+
19
+ class SimpleAgent:
20
+ """A tiny Chat Completions agent loop with workspace tools."""
21
+
22
+ def __init__(
23
+ self,
24
+ llm_client: Any,
25
+ model: str,
26
+ workspace: str = ".",
27
+ max_turns: int = 5,
28
+ command_approver: Callable[[str], bool] | None = None,
29
+ ) -> None:
30
+ self.llm_client = llm_client
31
+ self.model = model
32
+ self.max_turns = max_turns
33
+ self.command_approver = command_approver
34
+ self.tools = Tools(workspace)
35
+ self.conversation_history = ConversationHistory(INITIAL_PROMPT)
36
+
37
+ def respond_stream(self, user_input: str) -> Generator[TextDelta | ToolResult, None, None]:
38
+ self.conversation_history.append_user(user_input)
39
+
40
+ for _ in range(self.max_turns):
41
+ accumulator = ChatCompletionStreamAccumulator()
42
+ for chunk in self._create_llm_stream():
43
+ for event in accumulator.add_chunk(chunk):
44
+ yield event
45
+
46
+ message_dict = accumulator.assistant_message()
47
+ self.conversation_history.append(message_dict)
48
+
49
+ tool_calls = message_dict.get("tool_calls")
50
+ if not tool_calls:
51
+ return
52
+
53
+ for tool_call_dict in tool_calls:
54
+ result = self._handle_tool_call(tool_call_dict)
55
+ self.conversation_history.append_tool(tool_call_dict["id"], result)
56
+ yield ToolResult(
57
+ call_id=tool_call_dict["id"],
58
+ name=tool_call_dict["function"]["name"],
59
+ result=result,
60
+ ok=self._tool_result_ok(result),
61
+ )
62
+
63
+ final_text = "I stopped because the tool loop reached the max turn limit."
64
+ self.conversation_history.append_assistant(final_text)
65
+ yield TextDelta(final_text)
66
+
67
+ def _create_llm_stream(self) -> Iterable[Any]:
68
+ return self.llm_client.chat.completions.create(
69
+ model=self.model,
70
+ messages=self.conversation_history.to_messages(),
71
+ tools=self.tools.definitions(),
72
+ tool_choice="auto",
73
+ stream=True,
74
+ )
75
+
76
+ def _handle_tool_call(self, tool_call: ToolCallDict) -> str:
77
+ name = tool_call["function"]["name"]
78
+ try:
79
+ arguments = self._parse_tool_arguments(tool_call["function"]["arguments"])
80
+ except ValueError as exc:
81
+ return error(str(exc))
82
+ user_approval = self._approve_run_command(arguments)
83
+ if name == "run_command" and not user_approval:
84
+ return error("run_command was not approved")
85
+ return self.tools.run(name, arguments)
86
+
87
+ def _approve_run_command(self, arguments: dict[str, object]) -> bool:
88
+ command = arguments.get("command")
89
+ if not isinstance(command, str) or not command.strip():
90
+ # Empty or invalid commands are blocked by the tool implementation.
91
+ return True
92
+ if self.command_approver is None:
93
+ return False
94
+ return self.command_approver(command)
95
+
96
+ def _tool_result_ok(self, result: str) -> bool:
97
+ try:
98
+ parsed = json.loads(result)
99
+ except json.JSONDecodeError:
100
+ return False
101
+ if not isinstance(parsed, dict):
102
+ return False
103
+ return parsed.get("ok") is True
104
+
105
+ def _parse_tool_arguments(self, raw_arguments: str) -> dict[str, object]:
106
+ try:
107
+ arguments = json.loads(raw_arguments or "{}")
108
+ except json.JSONDecodeError as exc:
109
+ raise ValueError(f"invalid JSON arguments: {exc}")
110
+ if not isinstance(arguments, dict):
111
+ raise ValueError("tool arguments must be a JSON object")
112
+ return arguments
@@ -0,0 +1,43 @@
1
+ import os
2
+
3
+ from dotenv import load_dotenv
4
+
5
+ from simple_agent.agent_loop import SimpleAgent
6
+ from simple_agent.llm_client import (
7
+ DEEPSEEK_BASE_URL,
8
+ OPENAI_BASE_URL,
9
+ OpenAICompatibleClient,
10
+ )
11
+ from simple_agent.terminal_loop import approve_run_command, run_interactive_terminal
12
+
13
+
14
+ def main() -> None:
15
+ load_dotenv()
16
+
17
+ provider = os.environ.get("LLM_PROVIDER", "deepseek").lower()
18
+ timeout = float(os.environ.get("LLM_TIMEOUT", "60"))
19
+
20
+ if provider == "openai":
21
+ api_key = os.environ.get("OPENAI_API_KEY") or os.environ.get("LLM_API_KEY")
22
+ default_model = "gpt-5.4-mini"
23
+ default_base_url = OPENAI_BASE_URL
24
+ elif provider == "deepseek":
25
+ api_key = os.environ.get("DEEPSEEK_API_KEY") or os.environ.get("LLM_API_KEY")
26
+ default_model = "deepseek-v4-flash"
27
+ default_base_url = DEEPSEEK_BASE_URL
28
+ else:
29
+ raise ValueError("LLM_PROVIDER must be 'openai' or 'deepseek'")
30
+
31
+ llm_client = OpenAICompatibleClient(
32
+ api_key=api_key,
33
+ base_url=os.environ.get("LLM_BASE_URL", default_base_url),
34
+ timeout=timeout,
35
+ )
36
+ model = os.environ.get("LLM_MODEL", default_model)
37
+ agent = SimpleAgent(
38
+ llm_client=llm_client,
39
+ model=model,
40
+ command_approver=approve_run_command,
41
+ )
42
+
43
+ run_interactive_terminal(agent)
@@ -0,0 +1,59 @@
1
+ """Typed chat messages and conversation history storage."""
2
+
3
+ from collections.abc import Iterator
4
+ from copy import deepcopy
5
+ from typing import overload
6
+
7
+ from simple_agent.message_types import AssistantMessageDict, ChatMessage, ToolMessage, UserMessage
8
+
9
+
10
+ class ConversationHistory:
11
+ """Stores chat messages while exposing list-like read access."""
12
+ _messages: list[ChatMessage]
13
+
14
+ def __init__(self, system_prompt: str) -> None:
15
+ self._messages = [
16
+ {"role": "system", "content": system_prompt},
17
+ ]
18
+
19
+ def append(self, message: ChatMessage) -> None:
20
+ self._messages.append(message)
21
+
22
+ def append_user(self, content: str) -> None:
23
+ message: UserMessage = {"role": "user", "content": content}
24
+ self.append(message)
25
+
26
+ def append_assistant(self, content: str) -> None:
27
+ message: AssistantMessageDict = {"role": "assistant", "content": content}
28
+ self.append(message)
29
+
30
+ def append_tool(self, tool_call_id: str, content: str) -> None:
31
+ message: ToolMessage = {
32
+ "role": "tool",
33
+ "tool_call_id": tool_call_id,
34
+ "content": content,
35
+ }
36
+ self.append(message)
37
+
38
+ def replace(self, messages: list[ChatMessage]) -> None:
39
+ self._messages = messages
40
+
41
+ def to_messages(self) -> list[ChatMessage]:
42
+ return deepcopy(self._messages)
43
+
44
+ def __iter__(self) -> Iterator[ChatMessage]:
45
+ return iter(self._messages)
46
+
47
+ def __len__(self) -> int:
48
+ return len(self._messages)
49
+
50
+ @overload
51
+ def __getitem__(self, index: int) -> ChatMessage:
52
+ ...
53
+
54
+ @overload
55
+ def __getitem__(self, index: slice) -> list[ChatMessage]:
56
+ ...
57
+
58
+ def __getitem__(self, index: int | slice) -> ChatMessage | list[ChatMessage]:
59
+ return self._messages[index]
@@ -0,0 +1,49 @@
1
+ """OpenAI-compatible LLM clients."""
2
+
3
+ from openai import OpenAI
4
+
5
+
6
+ OPENAI_BASE_URL = "https://api.openai.com/v1"
7
+ DEEPSEEK_BASE_URL = "https://api.deepseek.com"
8
+
9
+
10
+ class LLMClientError(ValueError):
11
+ """Raised when LLM configuration is incomplete."""
12
+
13
+
14
+ class OpenAICompatibleClient:
15
+ """Small wrapper around the official OpenAI SDK client."""
16
+
17
+ base_url = OPENAI_BASE_URL
18
+
19
+ def __init__(
20
+ self,
21
+ api_key: str | None,
22
+ base_url: str | None = None,
23
+ timeout: float = 60.0,
24
+ ) -> None:
25
+ if not api_key:
26
+ raise LLMClientError("api_key is required to create an LLM client")
27
+
28
+ self.api_key = api_key
29
+ self.base_url = base_url or self.base_url
30
+ self.timeout = timeout
31
+ self.client = OpenAI(
32
+ base_url=self.base_url,
33
+ api_key=self.api_key,
34
+ timeout=self.timeout,
35
+ )
36
+
37
+ @property
38
+ def chat(self):
39
+ return self.client.chat
40
+
41
+
42
+ class OpenAIClient(OpenAICompatibleClient):
43
+ """OpenAI client."""
44
+
45
+
46
+ class DeepSeekClient(OpenAICompatibleClient):
47
+ """DeepSeek client using its OpenAI-compatible API."""
48
+
49
+ base_url = DEEPSEEK_BASE_URL
@@ -0,0 +1,39 @@
1
+ """TypedDict shapes for OpenAI-compatible chat messages."""
2
+
3
+ from typing import Literal, NotRequired, TypedDict
4
+
5
+
6
+ class SystemMessage(TypedDict):
7
+ role: Literal["system"]
8
+ content: str
9
+
10
+
11
+ class UserMessage(TypedDict):
12
+ role: Literal["user"]
13
+ content: str
14
+
15
+
16
+ class ToolMessage(TypedDict):
17
+ role: Literal["tool"]
18
+ tool_call_id: str
19
+ content: str
20
+
21
+
22
+ class FunctionCallDict(TypedDict):
23
+ name: str
24
+ arguments: str
25
+
26
+
27
+ class ToolCallDict(TypedDict):
28
+ id: str
29
+ type: Literal["function"]
30
+ function: FunctionCallDict
31
+
32
+
33
+ class AssistantMessageDict(TypedDict):
34
+ role: Literal["assistant"]
35
+ content: str | None
36
+ tool_calls: NotRequired[list[ToolCallDict]]
37
+
38
+
39
+ ChatMessage = SystemMessage | UserMessage | AssistantMessageDict | ToolMessage
@@ -0,0 +1,5 @@
1
+ INITIAL_PROMPT = """
2
+ You are a helpful educational coding assistant.
3
+ You may inspect and edit the workspace with tools before answering.
4
+ Use tools when file context would make your answer more accurate.
5
+ """
@@ -0,0 +1,96 @@
1
+ """Helpers for Chat Completions streaming responses."""
2
+
3
+ from dataclasses import dataclass
4
+
5
+ from openai.types.chat.chat_completion_chunk import (
6
+ ChatCompletionChunk,
7
+ ChoiceDeltaToolCall,
8
+ )
9
+
10
+ from simple_agent.message_types import AssistantMessageDict, ToolCallDict
11
+
12
+
13
+ @dataclass
14
+ class TextDelta:
15
+ """A streamed assistant text chunk."""
16
+
17
+ content: str
18
+
19
+
20
+ @dataclass
21
+ class ToolResult:
22
+ """A completed tool call event."""
23
+
24
+ call_id: str
25
+ name: str
26
+ result: str
27
+ ok: bool
28
+
29
+
30
+ def merge_tool_call_delta_into(tool_call: ToolCallDict, tool_call_delta: ChoiceDeltaToolCall) -> None:
31
+ # Each streamed delta may contain only one fragment of the final tool call.
32
+ call_id = tool_call_delta.id
33
+ if call_id:
34
+ tool_call["id"] = call_id
35
+
36
+ call_type = tool_call_delta.type
37
+ if call_type:
38
+ tool_call["type"] = call_type
39
+
40
+ function_delta = tool_call_delta.function
41
+ if function_delta is None:
42
+ return
43
+
44
+ name = function_delta.name
45
+ if name:
46
+ tool_call["function"]["name"] += name
47
+
48
+ arguments = function_delta.arguments
49
+ if arguments:
50
+ tool_call["function"]["arguments"] += arguments
51
+
52
+
53
+ class ChatCompletionStreamAccumulator:
54
+ """Build the final assistant message from streamed Chat Completions chunks."""
55
+
56
+ def __init__(self) -> None:
57
+ self.full_text = ""
58
+ self.tool_calls_by_index: dict[int, ToolCallDict] = {}
59
+
60
+ def add_chunk(self, chunk: ChatCompletionChunk) -> list[TextDelta]:
61
+ # OpenAI can send usage/bookkeeping chunks with choices=[].
62
+ choices = chunk.choices
63
+ if not choices:
64
+ return []
65
+
66
+ delta = choices[0].delta
67
+ events = []
68
+
69
+ content = delta.content
70
+ if content:
71
+ self.full_text += content
72
+ events.append(TextDelta(content))
73
+
74
+ for tool_call_delta in delta.tool_calls or []:
75
+ index = tool_call_delta.index
76
+ if index not in self.tool_calls_by_index:
77
+ self.tool_calls_by_index[index] = {
78
+ "id": "",
79
+ "type": "function",
80
+ "function": {"name": "", "arguments": ""},
81
+ }
82
+
83
+ tool_call = self.tool_calls_by_index[index]
84
+ # Later deltas usually add more argument text for this index.
85
+ merge_tool_call_delta_into(tool_call, tool_call_delta)
86
+
87
+ return events
88
+
89
+ def assistant_message(self) -> AssistantMessageDict:
90
+ message_dict: AssistantMessageDict = {"role": "assistant", "content": self.full_text or None}
91
+ if self.tool_calls_by_index:
92
+ message_dict["tool_calls"] = [
93
+ self.tool_calls_by_index[index]
94
+ for index in sorted(self.tool_calls_by_index)
95
+ ]
96
+ return message_dict