axio 0.3.5__tar.gz → 0.4.1__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.
- {axio-0.3.5 → axio-0.4.1}/PKG-INFO +4 -4
- {axio-0.3.5 → axio-0.4.1}/README.md +3 -3
- {axio-0.3.5 → axio-0.4.1}/pyproject.toml +1 -1
- {axio-0.3.5 → axio-0.4.1}/src/axio/agent.py +11 -3
- {axio-0.3.5 → axio-0.4.1}/src/axio/models.py +0 -9
- axio-0.4.1/src/axio/selector.py +10 -0
- axio-0.3.5/src/axio/selector.py +0 -121
- axio-0.3.5/tests/test_selector.py +0 -237
- {axio-0.3.5 → axio-0.4.1}/.github/workflows/publish.yml +0 -0
- {axio-0.3.5 → axio-0.4.1}/.github/workflows/tests.yml +0 -0
- {axio-0.3.5 → axio-0.4.1}/.gitignore +0 -0
- {axio-0.3.5 → axio-0.4.1}/LICENSE +0 -0
- {axio-0.3.5 → axio-0.4.1}/Makefile +0 -0
- {axio-0.3.5 → axio-0.4.1}/src/axio/__init__.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/src/axio/blocks.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/src/axio/context.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/src/axio/events.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/src/axio/exceptions.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/src/axio/messages.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/src/axio/permission.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/src/axio/py.typed +0 -0
- {axio-0.3.5 → axio-0.4.1}/src/axio/stream.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/src/axio/testing.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/src/axio/tool.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/src/axio/tool_args.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/src/axio/transport.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/src/axio/types.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/tests/conftest.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/tests/test_agent_branch.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/tests/test_agent_permission.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/tests/test_agent_run.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/tests/test_agent_stream.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/tests/test_agent_tools.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/tests/test_blocks.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/tests/test_context.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/tests/test_events.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/tests/test_exceptions.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/tests/test_permission.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/tests/test_stream.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/tests/test_tool.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/tests/test_tool_args.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/tests/test_transport.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/tests/test_types.py +0 -0
- {axio-0.3.5 → axio-0.4.1}/uv.lock +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: axio
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.4.1
|
|
4
4
|
Summary: Minimal, streaming-first, protocol-driven foundation for LLM-powered agents
|
|
5
5
|
Project-URL: Homepage, https://github.com/axio-agent/axio
|
|
6
6
|
Project-URL: Repository, https://github.com/axio-agent/axio
|
|
@@ -77,7 +77,7 @@ asyncio.run(main())
|
|
|
77
77
|
▼
|
|
78
78
|
┌─────────┐ stream() ┌─────────────────────┐
|
|
79
79
|
│ Agent │ ────────────▶ │ CompletionTransport │
|
|
80
|
-
│ loop │ ◀──────────── │ (
|
|
80
|
+
│ loop │ ◀──────────── │ (Anthropic, OpenAI, …) │
|
|
81
81
|
└─────────┘ StreamEvent └─────────────────────┘
|
|
82
82
|
│
|
|
83
83
|
│ tool_use?
|
|
@@ -202,8 +202,8 @@ my_guard = "my_package:MyGuard"
|
|
|
202
202
|
|
|
203
203
|
| Package | Purpose |
|
|
204
204
|
|---|---|
|
|
205
|
-
| [axio-transport-
|
|
206
|
-
| [axio-transport-
|
|
205
|
+
| [axio-transport-anthropic](https://github.com/axio-agent/axio-transport-anthropic) | Anthropic Claude transport |
|
|
206
|
+
| [axio-transport-openai](https://github.com/axio-agent/axio-transport-openai) | OpenAI-compatible transport (OpenAI, Nebius, OpenRouter, custom) |
|
|
207
207
|
| [axio-transport-codex](https://github.com/axio-agent/axio-transport-codex) | ChatGPT OAuth transport |
|
|
208
208
|
| [axio-tools-local](https://github.com/axio-agent/axio-tools-local) | Shell, file, Python tools |
|
|
209
209
|
| [axio-tools-mcp](https://github.com/axio-agent/axio-tools-mcp) | MCP server bridge |
|
|
@@ -64,7 +64,7 @@ asyncio.run(main())
|
|
|
64
64
|
▼
|
|
65
65
|
┌─────────┐ stream() ┌─────────────────────┐
|
|
66
66
|
│ Agent │ ────────────▶ │ CompletionTransport │
|
|
67
|
-
│ loop │ ◀──────────── │ (
|
|
67
|
+
│ loop │ ◀──────────── │ (Anthropic, OpenAI, …) │
|
|
68
68
|
└─────────┘ StreamEvent └─────────────────────┘
|
|
69
69
|
│
|
|
70
70
|
│ tool_use?
|
|
@@ -189,8 +189,8 @@ my_guard = "my_package:MyGuard"
|
|
|
189
189
|
|
|
190
190
|
| Package | Purpose |
|
|
191
191
|
|---|---|
|
|
192
|
-
| [axio-transport-
|
|
193
|
-
| [axio-transport-
|
|
192
|
+
| [axio-transport-anthropic](https://github.com/axio-agent/axio-transport-anthropic) | Anthropic Claude transport |
|
|
193
|
+
| [axio-transport-openai](https://github.com/axio-agent/axio-transport-openai) | OpenAI-compatible transport (OpenAI, Nebius, OpenRouter, custom) |
|
|
194
194
|
| [axio-transport-codex](https://github.com/axio-agent/axio-transport-codex) | ChatGPT OAuth transport |
|
|
195
195
|
| [axio-tools-local](https://github.com/axio-agent/axio-tools-local) | Shell, file, Python tools |
|
|
196
196
|
| [axio-tools-mcp](https://github.com/axio-agent/axio-tools-mcp) | MCP server bridge |
|
|
@@ -5,7 +5,7 @@ from __future__ import annotations
|
|
|
5
5
|
import asyncio
|
|
6
6
|
import json
|
|
7
7
|
import logging
|
|
8
|
-
from collections.abc import AsyncGenerator
|
|
8
|
+
from collections.abc import AsyncGenerator, Iterable
|
|
9
9
|
from dataclasses import dataclass, field
|
|
10
10
|
from typing import Any
|
|
11
11
|
|
|
@@ -22,6 +22,7 @@ from axio.events import (
|
|
|
22
22
|
ToolUseStart,
|
|
23
23
|
)
|
|
24
24
|
from axio.messages import Message
|
|
25
|
+
from axio.selector import ToolSelector
|
|
25
26
|
from axio.stream import AgentStream
|
|
26
27
|
from axio.tool import Tool
|
|
27
28
|
from axio.transport import CompletionTransport
|
|
@@ -35,6 +36,7 @@ class Agent:
|
|
|
35
36
|
system: str
|
|
36
37
|
tools: list[Tool]
|
|
37
38
|
transport: CompletionTransport
|
|
39
|
+
selector: ToolSelector | None = field(default=None)
|
|
38
40
|
max_iterations: int = field(default=50)
|
|
39
41
|
|
|
40
42
|
def run_stream(self, user_message: str, context: ContextStore) -> AgentStream:
|
|
@@ -120,6 +122,13 @@ class Agent:
|
|
|
120
122
|
blocks.append(ToolUseBlock(id=tid, name=info["name"], input=inp))
|
|
121
123
|
return blocks, malformed
|
|
122
124
|
|
|
125
|
+
async def _select_tools(self, history: list[Message], tools: list[Tool]) -> Iterable[Tool]:
|
|
126
|
+
if not tools:
|
|
127
|
+
return []
|
|
128
|
+
if not self.selector:
|
|
129
|
+
return tools
|
|
130
|
+
return await self.selector.select(history, tools)
|
|
131
|
+
|
|
123
132
|
async def _run_loop(self, user_message: str, context: ContextStore) -> AsyncGenerator[StreamEvent, None]:
|
|
124
133
|
total_usage = Usage(0, 0)
|
|
125
134
|
session_end_emitted = False
|
|
@@ -129,8 +138,7 @@ class Agent:
|
|
|
129
138
|
for iteration in range(1, self.max_iterations + 1):
|
|
130
139
|
history = await context.get_history()
|
|
131
140
|
logger.info("Iteration %d, history length=%d", iteration, len(history))
|
|
132
|
-
active_tools = self.tools
|
|
133
|
-
|
|
141
|
+
active_tools = list(await self._select_tools(history, self.tools))
|
|
134
142
|
content: list[TextBlock | ToolUseBlock] = []
|
|
135
143
|
pending: dict[str, dict[str, Any]] = {}
|
|
136
144
|
stop_reason = StopReason.end_turn
|
|
@@ -7,15 +7,6 @@ from dataclasses import dataclass
|
|
|
7
7
|
from enum import StrEnum
|
|
8
8
|
|
|
9
9
|
|
|
10
|
-
@dataclass(frozen=True, slots=True)
|
|
11
|
-
class TransportMeta:
|
|
12
|
-
"""Metadata a transport plugin declares about itself."""
|
|
13
|
-
|
|
14
|
-
label: str
|
|
15
|
-
api_key_env: str
|
|
16
|
-
role_defaults: dict[str, str]
|
|
17
|
-
|
|
18
|
-
|
|
19
10
|
class Capability(StrEnum):
|
|
20
11
|
text = "text"
|
|
21
12
|
vision = "vision"
|
|
@@ -0,0 +1,10 @@
|
|
|
1
|
+
from collections.abc import Iterable
|
|
2
|
+
from typing import Protocol, runtime_checkable
|
|
3
|
+
|
|
4
|
+
from axio.messages import Message
|
|
5
|
+
from axio.tool import Tool
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
@runtime_checkable
|
|
9
|
+
class ToolSelector(Protocol):
|
|
10
|
+
async def select(self, messages: Iterable[Message], tools: Iterable[Tool]) -> Iterable[Tool]: ...
|
axio-0.3.5/src/axio/selector.py
DELETED
|
@@ -1,121 +0,0 @@
|
|
|
1
|
-
"""Tool selector: choose relevant tools per query via embedding similarity."""
|
|
2
|
-
|
|
3
|
-
from __future__ import annotations
|
|
4
|
-
|
|
5
|
-
import math
|
|
6
|
-
from collections.abc import AsyncIterator
|
|
7
|
-
from typing import Protocol, runtime_checkable
|
|
8
|
-
|
|
9
|
-
from axio.blocks import TextBlock, ToolResultBlock
|
|
10
|
-
from axio.events import StreamEvent
|
|
11
|
-
from axio.messages import Message
|
|
12
|
-
from axio.tool import Tool
|
|
13
|
-
from axio.transport import CompletionTransport, EmbeddingTransport
|
|
14
|
-
|
|
15
|
-
|
|
16
|
-
@runtime_checkable
|
|
17
|
-
class ToolSelector(Protocol):
|
|
18
|
-
async def select(self, messages: list[Message], tools: list[Tool]) -> list[Tool]: ...
|
|
19
|
-
|
|
20
|
-
|
|
21
|
-
def _cosine_similarity(a: list[float], b: list[float]) -> float:
|
|
22
|
-
dot = sum(x * y for x, y in zip(a, b))
|
|
23
|
-
norm_a = math.sqrt(sum(x * x for x in a))
|
|
24
|
-
norm_b = math.sqrt(sum(x * x for x in b))
|
|
25
|
-
if norm_a == 0.0 or norm_b == 0.0:
|
|
26
|
-
return 0.0
|
|
27
|
-
return dot / (norm_a * norm_b)
|
|
28
|
-
|
|
29
|
-
|
|
30
|
-
def _extract_last_user_text(messages: list[Message]) -> str | None:
|
|
31
|
-
"""Return joined text from the last user message, or None if it's a tool-result iteration."""
|
|
32
|
-
for msg in reversed(messages):
|
|
33
|
-
if msg.role != "user":
|
|
34
|
-
continue
|
|
35
|
-
texts = [b.text for b in msg.content if isinstance(b, TextBlock)]
|
|
36
|
-
if texts:
|
|
37
|
-
return " ".join(texts)
|
|
38
|
-
# User message with only ToolResultBlocks → tool-result iteration
|
|
39
|
-
if any(isinstance(b, ToolResultBlock) for b in msg.content):
|
|
40
|
-
return None
|
|
41
|
-
return None
|
|
42
|
-
return None
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
class EmbeddingToolSelector:
|
|
46
|
-
def __init__(
|
|
47
|
-
self,
|
|
48
|
-
transport: EmbeddingTransport,
|
|
49
|
-
*,
|
|
50
|
-
top_k: int = 5,
|
|
51
|
-
pinned: frozenset[str] = frozenset(),
|
|
52
|
-
) -> None:
|
|
53
|
-
self._transport = transport
|
|
54
|
-
self._top_k = top_k
|
|
55
|
-
self._pinned = pinned
|
|
56
|
-
self._tool_embeddings: dict[str, list[float]] = {}
|
|
57
|
-
self._tool_descriptions: dict[str, str] = {}
|
|
58
|
-
|
|
59
|
-
async def _ensure_embeddings(self, tools: list[Tool]) -> None:
|
|
60
|
-
to_embed: list[tuple[str, str]] = []
|
|
61
|
-
for tool in tools:
|
|
62
|
-
cached_desc = self._tool_descriptions.get(tool.name)
|
|
63
|
-
if cached_desc is None or cached_desc != tool.description:
|
|
64
|
-
to_embed.append((tool.name, tool.description))
|
|
65
|
-
|
|
66
|
-
if not to_embed:
|
|
67
|
-
return
|
|
68
|
-
|
|
69
|
-
texts = [desc for _, desc in to_embed]
|
|
70
|
-
vectors = await self._transport.embed(texts)
|
|
71
|
-
for (name, desc), vec in zip(to_embed, vectors):
|
|
72
|
-
self._tool_embeddings[name] = vec
|
|
73
|
-
self._tool_descriptions[name] = desc
|
|
74
|
-
|
|
75
|
-
async def select(self, messages: list[Message], tools: list[Tool]) -> list[Tool]:
|
|
76
|
-
if len(tools) <= self._top_k:
|
|
77
|
-
return tools
|
|
78
|
-
|
|
79
|
-
query = _extract_last_user_text(messages)
|
|
80
|
-
if query is None:
|
|
81
|
-
return tools
|
|
82
|
-
|
|
83
|
-
await self._ensure_embeddings(tools)
|
|
84
|
-
|
|
85
|
-
query_vec = (await self._transport.embed([query]))[0]
|
|
86
|
-
|
|
87
|
-
pinned_tools: list[Tool] = []
|
|
88
|
-
scorable: list[tuple[float, Tool]] = []
|
|
89
|
-
|
|
90
|
-
for tool in tools:
|
|
91
|
-
if tool.name in self._pinned:
|
|
92
|
-
pinned_tools.append(tool)
|
|
93
|
-
elif tool.name in self._tool_embeddings:
|
|
94
|
-
score = _cosine_similarity(query_vec, self._tool_embeddings[tool.name])
|
|
95
|
-
scorable.append((score, tool))
|
|
96
|
-
|
|
97
|
-
scorable.sort(key=lambda x: x[0], reverse=True)
|
|
98
|
-
remaining_slots = max(0, self._top_k - len(pinned_tools))
|
|
99
|
-
selected = pinned_tools + [t for _, t in scorable[:remaining_slots]]
|
|
100
|
-
return selected
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
class ToolFilteringTransport:
|
|
104
|
-
"""CompletionTransport decorator that filters tools via a ToolSelector."""
|
|
105
|
-
|
|
106
|
-
def __init__(self, transport: CompletionTransport, selector: ToolSelector) -> None:
|
|
107
|
-
self._transport = transport
|
|
108
|
-
self._selector = selector
|
|
109
|
-
|
|
110
|
-
def stream(self, messages: list[Message], tools: list[Tool], system: str) -> AsyncIterator[StreamEvent]:
|
|
111
|
-
return self._filtered_stream(messages, tools, system)
|
|
112
|
-
|
|
113
|
-
async def _filtered_stream(
|
|
114
|
-
self,
|
|
115
|
-
messages: list[Message],
|
|
116
|
-
tools: list[Tool],
|
|
117
|
-
system: str,
|
|
118
|
-
) -> AsyncIterator[StreamEvent]:
|
|
119
|
-
selected = await self._selector.select(messages, tools)
|
|
120
|
-
async for event in self._transport.stream(messages, selected, system):
|
|
121
|
-
yield event
|
|
@@ -1,237 +0,0 @@
|
|
|
1
|
-
"""Tests for axon.selector — ToolSelector protocol, EmbeddingToolSelector, ToolFilteringTransport."""
|
|
2
|
-
|
|
3
|
-
from __future__ import annotations
|
|
4
|
-
|
|
5
|
-
import math
|
|
6
|
-
from collections.abc import AsyncGenerator
|
|
7
|
-
|
|
8
|
-
import pytest
|
|
9
|
-
|
|
10
|
-
from axio.blocks import TextBlock, ToolResultBlock
|
|
11
|
-
from axio.events import StreamEvent, TextDelta
|
|
12
|
-
from axio.messages import Message
|
|
13
|
-
from axio.selector import (
|
|
14
|
-
EmbeddingToolSelector,
|
|
15
|
-
ToolFilteringTransport,
|
|
16
|
-
ToolSelector,
|
|
17
|
-
_cosine_similarity,
|
|
18
|
-
_extract_last_user_text,
|
|
19
|
-
)
|
|
20
|
-
from axio.testing import StubTransport, make_text_response
|
|
21
|
-
from axio.tool import Tool, ToolHandler
|
|
22
|
-
from axio.transport import CompletionTransport
|
|
23
|
-
|
|
24
|
-
|
|
25
|
-
class DummyHandler(ToolHandler):
|
|
26
|
-
async def __call__(self) -> str:
|
|
27
|
-
return "ok"
|
|
28
|
-
|
|
29
|
-
|
|
30
|
-
def _make_tool(name: str, description: str = "") -> Tool:
|
|
31
|
-
return Tool(name=name, description=description or f"Tool {name}", handler=DummyHandler)
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
class StubEmbeddingTransport:
|
|
35
|
-
"""Returns deterministic vectors: hash-based unit vectors for reproducibility."""
|
|
36
|
-
|
|
37
|
-
def __init__(self) -> None:
|
|
38
|
-
self.call_count = 0
|
|
39
|
-
|
|
40
|
-
async def embed(self, texts: list[str]) -> list[list[float]]:
|
|
41
|
-
self.call_count += 1
|
|
42
|
-
return [self._vector_for(t) for t in texts]
|
|
43
|
-
|
|
44
|
-
@staticmethod
|
|
45
|
-
def _vector_for(text: str) -> list[float]:
|
|
46
|
-
"""Generate a deterministic 8-dim vector from text hash."""
|
|
47
|
-
h = hash(text) & 0xFFFFFFFF
|
|
48
|
-
raw = [float((h >> (i * 4)) & 0xF) - 7.5 for i in range(8)]
|
|
49
|
-
norm = math.sqrt(sum(x * x for x in raw))
|
|
50
|
-
if norm == 0:
|
|
51
|
-
return [0.0] * 8
|
|
52
|
-
return [x / norm for x in raw]
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
def test_protocol_conformance() -> None:
|
|
56
|
-
transport = StubEmbeddingTransport()
|
|
57
|
-
selector = EmbeddingToolSelector(transport=transport)
|
|
58
|
-
assert isinstance(selector, ToolSelector)
|
|
59
|
-
|
|
60
|
-
|
|
61
|
-
async def test_returns_all_when_few_tools() -> None:
|
|
62
|
-
transport = StubEmbeddingTransport()
|
|
63
|
-
selector = EmbeddingToolSelector(transport=transport, top_k=5)
|
|
64
|
-
tools = [_make_tool(f"t{i}") for i in range(3)]
|
|
65
|
-
messages = [Message(role="user", content=[TextBlock(text="hello")])]
|
|
66
|
-
|
|
67
|
-
result = await selector.select(messages, tools)
|
|
68
|
-
assert result == tools
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
async def test_selects_top_k() -> None:
|
|
72
|
-
transport = StubEmbeddingTransport()
|
|
73
|
-
selector = EmbeddingToolSelector(transport=transport, top_k=3)
|
|
74
|
-
tools = [_make_tool(f"tool_{i}", f"description {i}") for i in range(6)]
|
|
75
|
-
messages = [Message(role="user", content=[TextBlock(text="description 0")])]
|
|
76
|
-
|
|
77
|
-
result = await selector.select(messages, tools)
|
|
78
|
-
assert len(result) == 3
|
|
79
|
-
# All returned tools should be from the original list
|
|
80
|
-
for t in result:
|
|
81
|
-
assert t in tools
|
|
82
|
-
|
|
83
|
-
|
|
84
|
-
async def test_pinned_always_included() -> None:
|
|
85
|
-
transport = StubEmbeddingTransport()
|
|
86
|
-
selector = EmbeddingToolSelector(
|
|
87
|
-
transport=transport,
|
|
88
|
-
top_k=2,
|
|
89
|
-
pinned=frozenset({"pinned_tool"}),
|
|
90
|
-
)
|
|
91
|
-
tools = [
|
|
92
|
-
_make_tool("pinned_tool", "totally unrelated description xyz"),
|
|
93
|
-
_make_tool("tool_a", "alpha"),
|
|
94
|
-
_make_tool("tool_b", "beta"),
|
|
95
|
-
_make_tool("tool_c", "gamma"),
|
|
96
|
-
]
|
|
97
|
-
messages = [Message(role="user", content=[TextBlock(text="alpha beta")])]
|
|
98
|
-
|
|
99
|
-
result = await selector.select(messages, tools)
|
|
100
|
-
assert len(result) == 2
|
|
101
|
-
names = {t.name for t in result}
|
|
102
|
-
assert "pinned_tool" in names
|
|
103
|
-
|
|
104
|
-
|
|
105
|
-
async def test_returns_all_on_tool_result_iteration() -> None:
|
|
106
|
-
transport = StubEmbeddingTransport()
|
|
107
|
-
selector = EmbeddingToolSelector(transport=transport, top_k=2)
|
|
108
|
-
tools = [_make_tool(f"t{i}") for i in range(6)]
|
|
109
|
-
messages = [
|
|
110
|
-
Message(role="user", content=[TextBlock(text="hello")]),
|
|
111
|
-
Message(role="assistant", content=[TextBlock(text="running tool")]),
|
|
112
|
-
Message(role="user", content=[ToolResultBlock(tool_use_id="call_1", content="result")]),
|
|
113
|
-
]
|
|
114
|
-
|
|
115
|
-
result = await selector.select(messages, tools)
|
|
116
|
-
assert result == tools
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
def test_cosine_similarity_identical() -> None:
|
|
120
|
-
a = [1.0, 0.0, 0.0]
|
|
121
|
-
assert _cosine_similarity(a, a) == pytest.approx(1.0)
|
|
122
|
-
|
|
123
|
-
|
|
124
|
-
def test_cosine_similarity_orthogonal() -> None:
|
|
125
|
-
a = [1.0, 0.0]
|
|
126
|
-
b = [0.0, 1.0]
|
|
127
|
-
assert _cosine_similarity(a, b) == pytest.approx(0.0)
|
|
128
|
-
|
|
129
|
-
|
|
130
|
-
def test_cosine_similarity_opposite() -> None:
|
|
131
|
-
a = [1.0, 0.0]
|
|
132
|
-
b = [-1.0, 0.0]
|
|
133
|
-
assert _cosine_similarity(a, b) == pytest.approx(-1.0)
|
|
134
|
-
|
|
135
|
-
|
|
136
|
-
def test_cosine_similarity_zero_vector() -> None:
|
|
137
|
-
a = [0.0, 0.0]
|
|
138
|
-
b = [1.0, 0.0]
|
|
139
|
-
assert _cosine_similarity(a, b) == 0.0
|
|
140
|
-
|
|
141
|
-
|
|
142
|
-
async def test_caches_embeddings() -> None:
|
|
143
|
-
transport = StubEmbeddingTransport()
|
|
144
|
-
selector = EmbeddingToolSelector(transport=transport, top_k=2)
|
|
145
|
-
tools = [_make_tool(f"t{i}", f"desc {i}") for i in range(4)]
|
|
146
|
-
messages = [Message(role="user", content=[TextBlock(text="query")])]
|
|
147
|
-
|
|
148
|
-
await selector.select(messages, tools)
|
|
149
|
-
first_count = transport.call_count
|
|
150
|
-
|
|
151
|
-
await selector.select(messages, tools)
|
|
152
|
-
# Only the query embedding should be called again, not tool descriptions
|
|
153
|
-
assert transport.call_count == first_count + 1
|
|
154
|
-
|
|
155
|
-
|
|
156
|
-
def test_extract_last_user_text_with_text() -> None:
|
|
157
|
-
messages = [Message(role="user", content=[TextBlock(text="hello world")])]
|
|
158
|
-
assert _extract_last_user_text(messages) == "hello world"
|
|
159
|
-
|
|
160
|
-
|
|
161
|
-
def test_extract_last_user_text_tool_result() -> None:
|
|
162
|
-
messages = [
|
|
163
|
-
Message(role="user", content=[TextBlock(text="initial")]),
|
|
164
|
-
Message(role="assistant", content=[TextBlock(text="response")]),
|
|
165
|
-
Message(role="user", content=[ToolResultBlock(tool_use_id="call_1", content="ok")]),
|
|
166
|
-
]
|
|
167
|
-
assert _extract_last_user_text(messages) is None
|
|
168
|
-
|
|
169
|
-
|
|
170
|
-
def test_extract_last_user_text_empty() -> None:
|
|
171
|
-
assert _extract_last_user_text([]) is None
|
|
172
|
-
|
|
173
|
-
|
|
174
|
-
def test_extract_last_user_text_multiple_text_blocks() -> None:
|
|
175
|
-
messages = [Message(role="user", content=[TextBlock(text="hello"), TextBlock(text="world")])]
|
|
176
|
-
assert _extract_last_user_text(messages) == "hello world"
|
|
177
|
-
|
|
178
|
-
|
|
179
|
-
# --- ToolFilteringTransport tests ---
|
|
180
|
-
|
|
181
|
-
|
|
182
|
-
class StubSelector:
|
|
183
|
-
"""ToolSelector that keeps only tools whose names are in the allowed set."""
|
|
184
|
-
|
|
185
|
-
def __init__(self, allowed: set[str]) -> None:
|
|
186
|
-
self._allowed = allowed
|
|
187
|
-
|
|
188
|
-
async def select(self, messages: list[Message], tools: list[Tool]) -> list[Tool]:
|
|
189
|
-
return [t for t in tools if t.name in self._allowed]
|
|
190
|
-
|
|
191
|
-
|
|
192
|
-
def test_filtering_transport_protocol() -> None:
|
|
193
|
-
inner = StubTransport([make_text_response("ok")])
|
|
194
|
-
selector = StubSelector(allowed=set())
|
|
195
|
-
ft = ToolFilteringTransport(inner, selector)
|
|
196
|
-
assert isinstance(ft, CompletionTransport)
|
|
197
|
-
|
|
198
|
-
|
|
199
|
-
async def test_filtering_transport_filters_tools() -> None:
|
|
200
|
-
captured_tools: list[list[Tool]] = []
|
|
201
|
-
|
|
202
|
-
class CapturingTransport:
|
|
203
|
-
async def _gen(self, events: list[StreamEvent]) -> AsyncGenerator[StreamEvent, None]:
|
|
204
|
-
for e in events:
|
|
205
|
-
yield e
|
|
206
|
-
|
|
207
|
-
def stream(self, messages: list[Message], tools: list[Tool], system: str) -> AsyncGenerator[StreamEvent, None]:
|
|
208
|
-
captured_tools.append(tools)
|
|
209
|
-
return self._gen(make_text_response("ok"))
|
|
210
|
-
|
|
211
|
-
tools = [_make_tool("a"), _make_tool("b"), _make_tool("c")]
|
|
212
|
-
selector = StubSelector(allowed={"a", "c"})
|
|
213
|
-
ft = ToolFilteringTransport(CapturingTransport(), selector)
|
|
214
|
-
|
|
215
|
-
messages = [Message(role="user", content=[TextBlock(text="hello")])]
|
|
216
|
-
events: list[StreamEvent] = []
|
|
217
|
-
async for e in ft.stream(messages, tools, "sys"):
|
|
218
|
-
events.append(e)
|
|
219
|
-
|
|
220
|
-
assert len(captured_tools) == 1
|
|
221
|
-
names = [t.name for t in captured_tools[0]]
|
|
222
|
-
assert names == ["a", "c"]
|
|
223
|
-
|
|
224
|
-
|
|
225
|
-
async def test_filtering_transport_passes_events() -> None:
|
|
226
|
-
inner = StubTransport([make_text_response("hello")])
|
|
227
|
-
selector = StubSelector(allowed=set()) # allows nothing, but events still pass
|
|
228
|
-
ft = ToolFilteringTransport(inner, selector)
|
|
229
|
-
|
|
230
|
-
messages = [Message(role="user", content=[TextBlock(text="hi")])]
|
|
231
|
-
events: list[StreamEvent] = []
|
|
232
|
-
async for e in ft.stream(messages, [], "sys"):
|
|
233
|
-
events.append(e)
|
|
234
|
-
|
|
235
|
-
assert len(events) == 2
|
|
236
|
-
assert isinstance(events[0], TextDelta)
|
|
237
|
-
assert events[0].delta == "hello"
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|