初始化demo
This commit is contained in:
杜宇轩
@@ -0,0 +1,36 @@
|
||||
from typing import Optional, List
|
||||
from strands import Agent
|
||||
|
||||
from .conversation_manager import ByteWindowConversationManager
|
||||
from .hooks import default_hooks
|
||||
from .model import get_model
|
||||
from .prompts import build_system_prompt
|
||||
from .tools import get_tools
|
||||
from . import config
|
||||
|
||||
|
||||
|
||||
def create_agent(
|
||||
*,
|
||||
system_prompt: Optional[str] = None,
|
||||
messages: Optional[List[dict]] = None,
|
||||
) -> Agent:
|
||||
"""创建一个 Agent 实例，带有可选的系统提示和消息列表。
|
||||
|
||||
Args:
|
||||
system_prompt (Optional[str]): 系统提示，用于指导 Agent 的行为。
|
||||
messages (Optional[List[dict]]): 消息列表，每条消息是一个字典，包含角色和内容。
|
||||
|
||||
Returns:
|
||||
Agent: 创建的 Agent 实例。
|
||||
"""
|
||||
agent = Agent(
|
||||
model=get_model(),
|
||||
system_prompt=system_prompt,
|
||||
tools=get_tools(),
|
||||
conversation_manager=ByteWindowConversationManager(
|
||||
max_bytes=config.CONV_MAX_BYTES
|
||||
),
|
||||
hooks=default_hooks(),
|
||||
messages=messages or []
|
||||
)
|
||||
@@ -0,0 +1,40 @@
|
||||
import os
|
||||
|
||||
try:
|
||||
# Optional: load a local .env file if python-dotenv is installed
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
|
||||
# ===== Model backend selection =====
|
||||
MODEL_PROVIDER = os.getenv("MODEL_PROVIDER", "openai").strip().lower()
|
||||
MODEL_ID = os.getenv("MODEL_ID", "gpt-4o-mini").strip()
|
||||
|
||||
MAX_TOKENS = int(os.getenv("MAX_TOKENS", 1024))
|
||||
TEMPERATURE = float(os.getenv("TEMPERATURE", 0.7))
|
||||
|
||||
|
||||
# ===== Provider credentials / endpoints =====
|
||||
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "").strip()
|
||||
OPENAI_API_BASE_URL = os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1").strip()
|
||||
|
||||
ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY", "").strip()
|
||||
|
||||
OLLAMA_HOST = os.getenv("OLLAMA_HOST", "http://localhost:11434").strip()
|
||||
|
||||
AWS_REGION = os.getenv("AWS_REGION", "us-east-1").strip()
|
||||
|
||||
|
||||
# ===== Conversation / history =====
|
||||
CONV_MAX_BYTES = int(os.getenv("CONV_MAX_BYTES", 1024 * 1024)) # 1MB
|
||||
|
||||
|
||||
# ===== History backend: "json" (one file per session) or "sqlite" (single SQLite database) =====
|
||||
SESSION_BACKEND = os.getenv("SESSION_BACKEND", "json").strip().lower()
|
||||
SESSION_DIR = os.getenv("SESSION_DIR", "./sessions").strip()
|
||||
SQLITE_PATH = os.getenv("SQLITE_PATH", "./sessions/agent.db").strip()
|
||||
@@ -0,0 +1,44 @@
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from strands.agent.conversation_manager.sliding_window_conversation_manager import SlidingWindowConversationManager
|
||||
|
||||
DEFAULT_MAX_BYTES = 300000
|
||||
DEFAULT_MAX_WINDOW = 400
|
||||
|
||||
|
||||
class ByteWindowConversationManager(SlidingWindowConversationManager):
|
||||
def __init__(
|
||||
self,
|
||||
max_bytes: int = DEFAULT_MAX_BYTES,
|
||||
max_window: int = DEFAULT_MAX_WINDOW,
|
||||
**kwargs: Any
|
||||
):
|
||||
super().__init__(max_window=max_window, **kwargs)
|
||||
self.max_bytes = max_bytes
|
||||
|
||||
def _byte_window(self, messages: list) -> int:
|
||||
"""计算保留的消息数量"""
|
||||
total = 0
|
||||
keep = 0
|
||||
for m in reversed(messages):
|
||||
total += len(json.dumps(m, ensure_ascii=False, default=str).encode("utf-8"))
|
||||
if keep >= 1 and total > self.max_bytes:
|
||||
break
|
||||
keep += 1
|
||||
return max(keep, 2)
|
||||
|
||||
def apply_management(self, agent: "Any", **kwargs: Any) -> None:
|
||||
"""应用消息管理策略，确保消息总字节数不超过 max_bytes。"""
|
||||
messages = agent.messages
|
||||
eff = min(self.window_size, self._byte_window(messages))
|
||||
|
||||
if len(messages) <= eff:
|
||||
return
|
||||
|
||||
saved = self.window_size
|
||||
self.window_size = eff
|
||||
try:
|
||||
self.reduce_context(agent)
|
||||
finally:
|
||||
self.window_size = saved
|
||||
@@ -0,0 +1,3 @@
|
||||
from .logging_hook import ToolLoggingHook, default_hooks
|
||||
|
||||
__all__ = ["ToolLoggingHook", "default_hooks"]
|
||||
@@ -0,0 +1,23 @@
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from strands.hooks import BeforeToolCallEvent, HookProvicer, HookRegistry
|
||||
|
||||
logger = logging.getLogger("agent_demo.hooks")
|
||||
|
||||
class ToolLoggingHook(HookProvicer):
|
||||
"""A hook that logs tool calls and their results."""
|
||||
|
||||
def register(self, registry: HookRegistry, **kwargs: Any) -> None:
|
||||
registry.add_callback(BeforeToolCallEvent, self._on_before_tool)
|
||||
|
||||
def _on_before_tool(self, event: BeforeToolCallEvent) -> None:
|
||||
tool_use = getattr(event, "tool_use", None) or {}
|
||||
|
||||
name = tool_use.get("name") if isinstance(tool_use, dict) else None
|
||||
logger.info("[hook] about to call tool: %s", name or "<unknown>")
|
||||
|
||||
|
||||
def default_hooks() -> list:
|
||||
"""Hooks registered on every agent by default. Add consent/audit hooks here."""
|
||||
return [ToolLoggingHook()]
|
||||
@@ -0,0 +1,106 @@
|
||||
import threading
|
||||
from typing import Any
|
||||
|
||||
from sympy.physics.units import temperature
|
||||
|
||||
from . import config
|
||||
|
||||
|
||||
_shared_model = None
|
||||
_lock = threading.Lock()
|
||||
|
||||
|
||||
|
||||
|
||||
def _build_openai():
|
||||
from strands.models.openai import OpenAIModel
|
||||
|
||||
client_args:dict[str, Any] = {"api_key": config.OPENAI_API_KEY}
|
||||
if config.OPENAI_API_BASE_URL:
|
||||
client_args["base_url"] = config.OPENAI_API_BASE_URL
|
||||
|
||||
return OpenAIModel(
|
||||
client_args=client_args,
|
||||
model_id=config.MODEL_ID,
|
||||
params={"max_tokens": config.MAX_TOKENS, "temperature": config.TEMPERATURE},
|
||||
)
|
||||
|
||||
|
||||
def _build_anthropic():
|
||||
from strands.models.anthropic import AnthropicModel
|
||||
|
||||
return AnthropicModel(
|
||||
client_args={"api_key": config.ANTHROPIC_API_KEY},
|
||||
model_id=config.MODEL_ID,
|
||||
max_tokens=config.MAX_TOKENS,
|
||||
params={"temperature": config.TEMPERATURE},
|
||||
)
|
||||
|
||||
|
||||
def _build_litellm():
|
||||
from strands.models.litellm import LiteLLMModel
|
||||
|
||||
return LiteLLMModel(
|
||||
model_id=config.MODEL_ID,
|
||||
|
||||
params={"max_tokens": config.MAX_TOKENS, "temperature": config.TEMPERATURE},
|
||||
)
|
||||
|
||||
|
||||
def _build_ollama():
|
||||
from strands.models.ollama import OllamaModel
|
||||
|
||||
return OllamaModel(
|
||||
host=config.OLLAMA_HOST,
|
||||
model_id=config.MODEL_ID,
|
||||
params={"temperature": config.TEMPERATURE},
|
||||
)
|
||||
|
||||
|
||||
def _build_bedrock():
|
||||
from strands.models.bedrock import BedrockModel
|
||||
|
||||
return BedrockModel(
|
||||
model_id=config.MODEL_ID,
|
||||
region=config.AWS_REGION,
|
||||
max_tokens=config.MAX_TOKENS,
|
||||
temperature=config.TEMPERATURE,
|
||||
)
|
||||
|
||||
|
||||
_BUILDERS = {
|
||||
"openai": _build_openai,
|
||||
"anthropic": _build_anthropic,
|
||||
"litellm": _build_litellm,
|
||||
"ollama": _build_ollama,
|
||||
"bedrock": _build_bedrock,
|
||||
}
|
||||
|
||||
|
||||
def get_model():
|
||||
"""获取共享的模型实例，根据配置选择不同的模型提供者。
|
||||
|
||||
Returns:
|
||||
Any: 模型实例。
|
||||
"""
|
||||
global _shared_model
|
||||
if _shared_model is not None:
|
||||
return _shared_model
|
||||
|
||||
with _lock:
|
||||
if _shared_model is not None:
|
||||
return _shared_model
|
||||
|
||||
builder = _BUILDERS.get(config.MODEL_PROVIDER)
|
||||
if builder is None:
|
||||
raise ValueError(f"Unsupported model provider: {config.MODEL_PROVIDER}"
|
||||
f"Choose one of: {', '.join(_BUILDERS.keys())}")
|
||||
_shared_model = builder()
|
||||
return _shared_model
|
||||
|
||||
|
||||
def reset_model() -> None:
|
||||
"""重置共享的模型实例，通常用于测试或重新加载配置。"""
|
||||
global _shared_model
|
||||
with _lock:
|
||||
_shared_model = None
|
||||
@@ -0,0 +1,23 @@
|
||||
from .base import Session, SessionStore
|
||||
from .json_store import JSONSessionStore
|
||||
from .sqlite_store import SQLiteSessionStore
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Session",
|
||||
"SessionStore",
|
||||
"JSONSessionStore",
|
||||
"SQLiteSessionStore",
|
||||
"get_session_store",
|
||||
]
|
||||
|
||||
|
||||
def get_session_store() -> SessionStore:
|
||||
from .. import config
|
||||
backend = config.SESSION_BACKEND
|
||||
if backend == "sqlite":
|
||||
return SQLiteSessionStore(config.SQLITE_PATH)
|
||||
if backend == "json":
|
||||
return JSONSessionStore(config.SESSION_DIR)
|
||||
raise ValueError(f"Unsupported SESSION_BACKEND: {backend}")
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
import re
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List
|
||||
|
||||
SAFE_ID = re.compile(r"^[a-zA-Z0-9._-]")
|
||||
|
||||
|
||||
def sanitize_id(value: str, fallback:str) -> str:
|
||||
"""Sanitize a string to be used as an ID by removing unsafe characters."""
|
||||
return SAFE_ID.sub("_", value) or fallback
|
||||
|
||||
|
||||
@dataclass
|
||||
class Session:
|
||||
"""Conversasion history for one (user_id, session_id) pair.
|
||||
``message`` is in Strands / Bedrock-converse format (provider-neutral::
|
||||
{"role": "user"|"assistant"|"system",
|
||||
"content": [{"text": ...} | {"toolUse": ...} | {"toolResult": ...} ]})
|
||||
"""
|
||||
session_id: str
|
||||
user_id: str
|
||||
messages: List[dict] = field(default_factory=list)
|
||||
|
||||
|
||||
class SessionStore(ABC):
|
||||
"""Abstract base class for session storage backends."""
|
||||
|
||||
@abstractmethod
|
||||
def load(self, user_id: str, session_id: str) -> Session:
|
||||
"""Return prior history for the key (empty Session if none exists yet)."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def save(self, session: Session) -> None:
|
||||
"""Save a session to the store."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def clear(self, user_id: str, session_id: str) -> None:
|
||||
"""Clear a session from the store."""
|
||||
pass
|
||||
@@ -0,0 +1,53 @@
|
||||
import json
|
||||
import os
|
||||
|
||||
from .base import Session, SessionStore, sanitize_id
|
||||
|
||||
|
||||
class JSONSessionStore(SessionStore):
|
||||
"""A simple JSON file-based session store."""
|
||||
|
||||
def __init__(self, base_dir: str):
|
||||
self.base_dir = base_dir
|
||||
os.makedirs(base_dir, exist_ok=True)
|
||||
|
||||
def _path(self, user_id: str, session_id: str) -> str:
|
||||
return os.path.join(
|
||||
self.base_dir,
|
||||
f"{sanitize_id(user_id, 'anon')}__{sanitize_id(session_id, 'default')}.json")
|
||||
|
||||
def load(self, user_id: str, session_id: str) -> Session:
|
||||
path = self._path(user_id, session_id)
|
||||
messages = []
|
||||
if os.path.exists(path):
|
||||
try:
|
||||
with open(path, 'r', encoding='utf-8') as f:
|
||||
messages = json.load(f).get('messages', [])
|
||||
|
||||
except (json.JSONDecodeError, OSError):
|
||||
messages = []
|
||||
|
||||
return Session(session_id=session_id, user_id=user_id, messages=messages)
|
||||
|
||||
|
||||
def save(self, session: Session) -> None:
|
||||
path = self._path(session.user_id, session.session_id)
|
||||
tmp = path + ".tmp"
|
||||
|
||||
with open(path, 'w', encoding='utf-8') as f:
|
||||
json.dump(
|
||||
{
|
||||
"user_id": session.user_id,
|
||||
"session_id": session.session_id,
|
||||
"messages": session.messages,
|
||||
},
|
||||
f,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
os.replace(tmp, path)
|
||||
|
||||
def clear(self, user_id: str, session_id: str) -> None:
|
||||
path = self._path(user_id, session_id)
|
||||
if os.path.exists(path):
|
||||
os.remove(path)
|
||||
@@ -0,0 +1,74 @@
|
||||
import json
|
||||
import sqlite3
|
||||
from typing import Optional
|
||||
|
||||
from .base import Session, SessionStore
|
||||
|
||||
_SCHEMA = """
|
||||
CREATE TABLE IF NOT EXISTS sessions (
|
||||
user_id TEXT NOT NULL,
|
||||
session_id TEXT NOT NULL,
|
||||
messages TEXT NOT NULL, -- JSON-encoded message list
|
||||
updated_at REAL NOT NULL,
|
||||
PRIMARY KEY (user_id, session_id)
|
||||
);
|
||||
"""
|
||||
|
||||
|
||||
class SQLiteSessionStore(SessionStore):
|
||||
"""A simple SQLite-based session store."""
|
||||
|
||||
def __init__(self, db_path: str):
|
||||
self.db_path = db_path
|
||||
self._init_schema()
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
conn = sqlite3.connect(self.db_path, check_same_thread=False)
|
||||
conn.row_factory = sqlite3.Row
|
||||
return conn
|
||||
|
||||
def _init_schema(self):
|
||||
with self._connect() as conn:
|
||||
conn.execute(_SCHEMA)
|
||||
|
||||
def load(self, user_id: str, session_id: str) -> Session:
|
||||
with self._connect() as conn:
|
||||
row: Optional[sqlite3.Row] = conn.execute(
|
||||
"SELECT messages FROM agent_session WHERE user_id = ? AND session_id = ?",
|
||||
(user_id, session_id),
|
||||
).fetchone()
|
||||
messages = []
|
||||
if row is not None:
|
||||
try:
|
||||
messages = json.loads(row["messages"])
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
messages = []
|
||||
return Session(session_id=session_id, user_id=user_id, messages=messages)
|
||||
|
||||
def save(self, session: Session) -> None:
|
||||
import time
|
||||
|
||||
with self._connect() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO agent_session (user_id, session_id, messages, updated_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
ON CONFLICT(user_id, session_id)
|
||||
DO UPDATE SET messages=excluded.messages, updated_at=excluded.updated_at
|
||||
""",
|
||||
(
|
||||
session.user_id,
|
||||
session.session_id,
|
||||
json.dumps(session.messages, ensure_ascii=False),
|
||||
time.time()
|
||||
),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def clear(self, user_id: str, session_id: str) -> None:
|
||||
with self._connect() as conn:
|
||||
conn.execute(
|
||||
"DELETE FROM agent_session WHERE user_id = ? AND session_id = ?",
|
||||
(user_id, session_id),
|
||||
)
|
||||
conn.commit()
|
||||
@@ -0,0 +1,91 @@
|
||||
import asyncio
|
||||
from dataclasses import dataclass
|
||||
from typing import AsyncIterator, Callable, Optional
|
||||
|
||||
from .agent_factory import create_agent
|
||||
from .persistence import Session, SessionStore, get_session_store
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamEvent:
|
||||
"""Represents a streaming event from the agent."""
|
||||
|
||||
kind: str
|
||||
text: str = ""
|
||||
tool_name: str = ""
|
||||
|
||||
|
||||
async def stream_turn(agent, user_input: str) -> AsyncIterator[StreamEvent]:
|
||||
"""Streams the agent's response to a user input in real-time.
|
||||
|
||||
Yields token/tool events as they arrive, then a final "done" event
|
||||
"""
|
||||
last_tool_id = None
|
||||
|
||||
async for event in agent.stream_async(user_input):
|
||||
if event.get("data"):
|
||||
yield StreamEvent(kind="token", text=event["data"])
|
||||
elif event.get("current_tool_use"):
|
||||
tu = event["current_tool_use"]
|
||||
tid = tu.get("toolUsedId") if isinstance(tu, dict) else None
|
||||
# De-dupe: Strands re-emits current_tool_use on every token while the
|
||||
# tool call is being assembled; surface one frame per distinct call.
|
||||
|
||||
if tid and tid != last_tool_id:
|
||||
last_tool_id = tid
|
||||
yield StreamEvent(kind="tool", tool_name=tu.get("toolName", ""))
|
||||
yield StreamEvent(kind="done")
|
||||
|
||||
|
||||
async def run_turn(
|
||||
user_id: str,
|
||||
session_id_str,
|
||||
user_input: str,
|
||||
*,
|
||||
store: Optional[SessionStore] = None,
|
||||
on_event: Optional[Callable[[StreamEvent], None]] = None,
|
||||
system_prompt: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Runs one full turn end-to-end with history persistence.
|
||||
|
||||
1. Load prior history for (user_id, session_id).
|
||||
2. Build a fresh agent seeded with that history.
|
||||
3. Stream the turn, forwarding events to "on_event" if provided.
|
||||
4. Persist the agent's updated messages back to the store.
|
||||
|
||||
Returns the full assistant reply text.
|
||||
"""
|
||||
store = store or get_session_store()
|
||||
|
||||
# 1. Load prior history
|
||||
session = store.load(user_id, session_id)
|
||||
|
||||
# 2. SEED a fresh agent with the prior messages
|
||||
agent = create_agent(
|
||||
system_prompt=system_prompt,
|
||||
messages=session.messages
|
||||
)
|
||||
|
||||
# 3. Stream the turn
|
||||
reply_parts: list[str] = []
|
||||
async for event in stream_turn(agent, user_input):
|
||||
if event.kind == "token":
|
||||
reply_parts.append(event.text)
|
||||
if on_event is not None:
|
||||
on_event(event)
|
||||
|
||||
# 4. PERSIST updated history. Strands has appended this turn's user message,
|
||||
# any tool-use/tool-result messages, and the assistant reply to
|
||||
# agent.messages - persisting that list is the whole history contract.
|
||||
store.save(Session(session_id=session_id, user_id=user_id, messages=list(agent.messages)))
|
||||
return "".join(reply_parts)
|
||||
|
||||
|
||||
def run_turn_sync(
|
||||
user_id: str,
|
||||
session_id: str,
|
||||
user_input: str,
|
||||
**kwargs
|
||||
) -> str:
|
||||
"""Blocking convenience wrapper around :func: run_turn."""
|
||||
return asyncio.run(run_turn(user_id, session_id, user_input, **kwargs))
|
||||
@@ -0,0 +1 @@
|
||||
s
|
||||
Reference in New Issue
Block a user