ai-agent-book 精选快照(<2MB 代码与文档,来自 github.com/bojieli/ai-agent-book)
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,87 @@
|
||||
from typing import AsyncGenerator
|
||||
from aworld.cmd.utils.agent_server import AgentServer
|
||||
from aworld.output.ui.base import AworldUI
|
||||
from aworld.output.workspace import WorkSpace
|
||||
from aworld.cmd.data_model import (
|
||||
BaseAWorldAgent,
|
||||
ChatCompletionChoice,
|
||||
ChatCompletionMessage,
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionResponse,
|
||||
)
|
||||
from .agent_ui_parser import AWorldWebAgentUI
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from dotenv import load_dotenv
|
||||
import traceback
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def stream_run(request: ChatCompletionRequest, agent_server: AgentServer):
|
||||
if not request.session_id:
|
||||
request.session_id = str(uuid.uuid4())
|
||||
if not request.query_id:
|
||||
request.query_id = str(uuid.uuid4())
|
||||
if request.messages and request.messages[-1].trace_id is None:
|
||||
request.messages[-1].trace_id = request.trace_id
|
||||
|
||||
logger.info(f"Stream run agent: request={request.model_dump_json()}")
|
||||
agent = agent_server.get_agent(request.model)
|
||||
instance: BaseAWorldAgent = agent.instance
|
||||
env_file = os.path.join(agent.path, ".env")
|
||||
if os.path.exists(env_file):
|
||||
logger.info(f"Loading environment variables from {env_file}")
|
||||
load_dotenv(env_file, override=True, verbose=True)
|
||||
|
||||
final_response: str = ""
|
||||
|
||||
def build_response(delta_content: str):
|
||||
nonlocal final_response
|
||||
final_response += delta_content
|
||||
logger.info(f"Agent {agent.name} response: {delta_content}")
|
||||
return ChatCompletionResponse(
|
||||
choices=[
|
||||
ChatCompletionChoice(
|
||||
index=0,
|
||||
delta=ChatCompletionMessage(
|
||||
role="assistant",
|
||||
content=delta_content,
|
||||
trace_id=request.trace_id,
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
rich_ui = AWorldWebAgentUI(
|
||||
session_id=request.session_id,
|
||||
workspace=WorkSpace.from_local_storages(
|
||||
workspace_id=request.session_id,
|
||||
),
|
||||
)
|
||||
|
||||
await agent_server.on_chat_completion_request(request)
|
||||
try:
|
||||
async for output in instance.run(request=request):
|
||||
try:
|
||||
logger.info(f"Agent {agent.name} output: {output}")
|
||||
|
||||
if isinstance(output, str):
|
||||
yield build_response(output)
|
||||
else:
|
||||
res = await AworldUI.parse_output(output, rich_ui)
|
||||
for item in res if isinstance(res, list) else [res]:
|
||||
if isinstance(item, AsyncGenerator):
|
||||
async for sub_item in item:
|
||||
yield build_response(sub_item)
|
||||
else:
|
||||
yield build_response(item)
|
||||
except:
|
||||
logger.error(
|
||||
f"Agent {agent.name} output error! output={output}, error={traceback.format_exc()}"
|
||||
)
|
||||
except:
|
||||
logger.error(f"Agent {agent.name} error: {traceback.format_exc()}")
|
||||
finally:
|
||||
await agent_server.on_chat_completion_end(request, final_response)
|
||||
@@ -0,0 +1,112 @@
|
||||
import os
|
||||
import importlib
|
||||
import subprocess
|
||||
import sys
|
||||
import traceback
|
||||
import logging
|
||||
from typing import List, Dict
|
||||
from aworld.cmd.data_model import AgentModel
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_agent_cache: Dict[str, AgentModel] = {}
|
||||
|
||||
|
||||
def list_agents(server_dir: str) -> Dict[str, AgentModel]:
|
||||
"""
|
||||
List all cached agents
|
||||
|
||||
Returns:
|
||||
Dict[str, AgentModel]: The map of agent models
|
||||
"""
|
||||
if len(_agent_cache) == 0:
|
||||
for m in _list_agents(server_dir):
|
||||
_agent_cache[m.id] = m
|
||||
return _agent_cache
|
||||
|
||||
|
||||
def _list_agents(server_dir: str) -> List[AgentModel]:
|
||||
agents_dir = os.path.join(server_dir, "agent_deploy")
|
||||
|
||||
if not os.path.exists(agents_dir):
|
||||
logger.warning(f"Agents directory {agents_dir} does not exist")
|
||||
return []
|
||||
|
||||
if agents_dir not in sys.path:
|
||||
sys.path.append(agents_dir)
|
||||
|
||||
agents = []
|
||||
for agent_id in os.listdir(agents_dir):
|
||||
if agent_id.startswith("_"):
|
||||
continue
|
||||
try:
|
||||
agent_path = os.path.join(agents_dir, agent_id)
|
||||
if os.path.isdir(agent_path):
|
||||
requirements_file = os.path.join(agent_path, "requirements.txt")
|
||||
if os.path.exists(requirements_file):
|
||||
p = subprocess.Popen(
|
||||
["pip", "install", "-U", "-r", requirements_file],
|
||||
cwd=agent_path,
|
||||
)
|
||||
p.wait()
|
||||
if p.returncode != 0:
|
||||
logger.error(
|
||||
f"Error installing requirements for agent {agent_id}, path {agent_path}"
|
||||
)
|
||||
continue
|
||||
|
||||
agent_file = os.path.join(agent_path, "agent.py")
|
||||
if os.path.exists(agent_file):
|
||||
try:
|
||||
instance = _get_agent_instance(agent_id)
|
||||
if hasattr(instance, "name"):
|
||||
name = instance.name()
|
||||
else:
|
||||
name = agent_id
|
||||
if hasattr(instance, "description"):
|
||||
description = instance.description()
|
||||
else:
|
||||
description = ""
|
||||
agent_model = AgentModel(
|
||||
id=agent_id,
|
||||
name=name,
|
||||
description=description,
|
||||
path=agent_path,
|
||||
instance=instance,
|
||||
)
|
||||
|
||||
agents.append(agent_model)
|
||||
logger.info(
|
||||
f"Loaded agent {agent_id} successfully, path {agent_path}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error loading agent {agent_id}: {traceback.format_exc()}"
|
||||
)
|
||||
continue
|
||||
else:
|
||||
logger.warning(f"Agent {agent_id} does not have agent.py file")
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error loading agent {agent_id}, path {agent_path} : {traceback.format_exc()}"
|
||||
)
|
||||
continue
|
||||
|
||||
return agents
|
||||
|
||||
|
||||
def _get_agent_instance(agent_name):
|
||||
try:
|
||||
agent_module = importlib.import_module(
|
||||
name=f"{agent_name}.agent",
|
||||
)
|
||||
except Exception as e:
|
||||
msg = f"Error loading agent {agent_name}, cwd:{os.getcwd()}, sys.path:{sys.path}: {traceback.format_exc()}"
|
||||
logger.error(msg)
|
||||
raise Exception(msg)
|
||||
|
||||
if hasattr(agent_module, "AWorldAgent"):
|
||||
agent = agent_module.AWorldAgent()
|
||||
return agent
|
||||
else:
|
||||
raise Exception(f"Agent {agent_name} does not have AWorldAgent class")
|
||||
@@ -0,0 +1,116 @@
|
||||
from abc import abstractmethod
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from typing import Dict, List
|
||||
import os
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from aworld import trace
|
||||
from aworld.cmd.data_model import (
|
||||
AgentModel,
|
||||
ChatCompletionMessage,
|
||||
ChatCompletionRequest,
|
||||
)
|
||||
from aworld.session.base_session_service import BaseSessionService
|
||||
from aworld.session.simple_session_service import SimpleSessionService
|
||||
from . import agent_loader
|
||||
from aworld.trace.config import ObservabilityConfig
|
||||
from aworld.trace.opentelemetry.memory_storage import InMemoryWithPersistStorage
|
||||
|
||||
|
||||
# bugfix for tracer exception
|
||||
trace.configure(ObservabilityConfig(trace_storage=(InMemoryWithPersistStorage())))
|
||||
|
||||
|
||||
class ChatCallBack:
|
||||
@abstractmethod
|
||||
async def on_chat_completion_request(
|
||||
self, server: "AgentServer", request: ChatCompletionRequest
|
||||
):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def on_chat_completion_end(
|
||||
self, server: "AgentServer", request: ChatCompletionRequest, final_response: str
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
class SessionChatCallBack(ChatCallBack):
|
||||
|
||||
async def on_chat_completion_request(
|
||||
self, server: "AgentServer", request: ChatCompletionRequest
|
||||
):
|
||||
await server.get_session_service().append_messages(
|
||||
request.user_id,
|
||||
request.session_id,
|
||||
request.messages[-1:],
|
||||
)
|
||||
|
||||
async def on_chat_completion_end(
|
||||
self, server: "AgentServer", request: ChatCompletionRequest, final_response: str
|
||||
):
|
||||
await server.get_session_service().append_messages(
|
||||
request.user_id,
|
||||
request.session_id,
|
||||
[
|
||||
ChatCompletionMessage(
|
||||
role="assistant",
|
||||
content=final_response,
|
||||
trace_id=request.trace_id,
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class AgentServer:
|
||||
server_id: str
|
||||
server_name: str
|
||||
server_dir: str
|
||||
session_service: BaseSessionService = None
|
||||
agent_instances: Dict[str, AgentModel] = {}
|
||||
chat_call_backs: List[ChatCallBack] = [SessionChatCallBack()]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
server_id: str,
|
||||
server_name: str,
|
||||
server_dir: str = os.getcwd(),
|
||||
session_service: BaseSessionService = SimpleSessionService(),
|
||||
):
|
||||
"""
|
||||
Initialize AgentServer
|
||||
"""
|
||||
self.server_id = server_id
|
||||
self.server_name = server_name
|
||||
self.server_dir = server_dir
|
||||
self.session_service = session_service
|
||||
# Load server global env
|
||||
load_dotenv(Path(self.server_dir) / ".env", override=True, verbose=True)
|
||||
# Load agent instances
|
||||
self.agent_instances = agent_loader.list_agents(self.server_dir)
|
||||
|
||||
def list_agents(self) -> Dict[str, AgentModel]:
|
||||
return self.agent_instances
|
||||
|
||||
def get_agent(self, agent_id: str) -> AgentModel:
|
||||
return self.agent_instances.get(agent_id)
|
||||
|
||||
def get_session_service(self) -> BaseSessionService:
|
||||
return self.session_service
|
||||
|
||||
async def on_chat_completion_request(self, request: ChatCompletionRequest):
|
||||
tasks = []
|
||||
for chat_call_back in self.chat_call_backs:
|
||||
tasks.append(chat_call_back.on_chat_completion_request(self, request))
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
async def on_chat_completion_end(
|
||||
self, request: ChatCompletionRequest, final_response: str
|
||||
):
|
||||
tasks = []
|
||||
for chat_call_back in self.chat_call_backs:
|
||||
tasks.append(
|
||||
chat_call_back.on_chat_completion_end(self, request, final_response)
|
||||
)
|
||||
await asyncio.gather(*tasks)
|
||||
@@ -0,0 +1,253 @@
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
import uuid
|
||||
|
||||
from pydantic import Field, BaseModel, ConfigDict
|
||||
|
||||
from aworld.output import (
|
||||
MessageOutput,
|
||||
AworldUI,
|
||||
Output,
|
||||
WorkSpace,
|
||||
)
|
||||
from aworld.output.artifact import Artifact, ArtifactType
|
||||
from aworld.output.base import StepOutput, ToolResultOutput
|
||||
from aworld.output.utils import consume_content
|
||||
from abc import ABC, abstractmethod
|
||||
from typing_extensions import override
|
||||
|
||||
|
||||
class ToolCard(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
tool_type: str = Field(None, description="tool type")
|
||||
tool_name: str = Field(None, description="tool name")
|
||||
function_name: str = Field(None, description="function name")
|
||||
tool_call_id: str = Field(None, description="tool call id")
|
||||
arguments: str = Field(None, description="arguments")
|
||||
results: str = Field(None, description="results")
|
||||
card_type: str = Field(None, description="card type")
|
||||
card_data: dict = Field(None, description="card data")
|
||||
artifacts: list = Field(default_factory=list, description="artifacts")
|
||||
|
||||
@staticmethod
|
||||
def from_tool_result(output: ToolResultOutput) -> "ToolCard":
|
||||
return ToolCard(
|
||||
tool_type=output.tool_type,
|
||||
tool_name=output.tool_name,
|
||||
function_name=output.origin_tool_call.function.name,
|
||||
tool_call_id=output.origin_tool_call.id,
|
||||
arguments=output.origin_tool_call.function.arguments,
|
||||
results=output.data,
|
||||
artifacts=[],
|
||||
)
|
||||
|
||||
|
||||
class BaseToolResultParser(ABC):
|
||||
|
||||
def __init__(self, tool_name: str = None):
|
||||
self.tool_name = tool_name or self.__class__.__name__
|
||||
|
||||
@abstractmethod
|
||||
async def parse(self, output: ToolResultOutput, workspace: WorkSpace):
|
||||
pass
|
||||
|
||||
|
||||
class DefaultToolResultParser(BaseToolResultParser):
|
||||
|
||||
@override
|
||||
async def parse(self, output: ToolResultOutput, workspace: WorkSpace):
|
||||
tool_card = ToolCard.from_tool_result(output)
|
||||
|
||||
tool_card.card_type = "tool_call_card_default"
|
||||
|
||||
# screenshots
|
||||
if (
|
||||
output.metadata.get("screenshots")
|
||||
and isinstance(output.metadata.get("screenshots"), list)
|
||||
and len(output.metadata.get("screenshots")) > 0
|
||||
):
|
||||
for _, screenshot in enumerate(output.metadata.get("screenshots")):
|
||||
image_artifact = Artifact(
|
||||
artifact_id=str(uuid.uuid4()),
|
||||
artifact_type=ArtifactType.IMAGE,
|
||||
content=screenshot.get("ossPath"),
|
||||
)
|
||||
await workspace.add_artifact(image_artifact)
|
||||
tool_card.artifacts.append(
|
||||
{
|
||||
"artifact_type": image_artifact.artifact_type.value,
|
||||
"artifact_id": image_artifact.artifact_id,
|
||||
}
|
||||
)
|
||||
|
||||
return f"""\
|
||||
\n\n**🔧 Tool: {tool_card.tool_name}#{tool_card.function_name}**\n\n
|
||||
```tool_card
|
||||
{json.dumps(tool_card.model_dump(), ensure_ascii=False, indent=2)}
|
||||
```\n
|
||||
"""
|
||||
|
||||
|
||||
class SearchToolResultParser(BaseToolResultParser):
|
||||
|
||||
@override
|
||||
async def parse(self, output: ToolResultOutput, workspace: WorkSpace):
|
||||
tool_card = ToolCard.from_tool_result(output)
|
||||
|
||||
query = ""
|
||||
try:
|
||||
args = json.loads(tool_card.arguments)
|
||||
query = args.get("query")
|
||||
# aworld search server
|
||||
if not query:
|
||||
query = args.get("query_list")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
result_items = []
|
||||
try:
|
||||
result_items = json.loads(tool_card.results)
|
||||
# aworld search server return url, not link
|
||||
if result_items and isinstance(result_items, list):
|
||||
for item in result_items:
|
||||
if not item.get("link", None) and item.get("url", None):
|
||||
item["link"] = item.get("url")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if len(result_items) > 0:
|
||||
tool_card.results = ""
|
||||
|
||||
tool_card.card_type = "tool_call_card_link_list"
|
||||
tool_card.card_data = {
|
||||
"title": "🔎 Google Search",
|
||||
"query": query,
|
||||
"search_items": result_items,
|
||||
}
|
||||
|
||||
artifact_id = str(uuid.uuid4())
|
||||
await workspace.create_artifact(
|
||||
artifact_type=ArtifactType.WEB_PAGES,
|
||||
artifact_id=artifact_id,
|
||||
content=result_items,
|
||||
metadata={
|
||||
"query": query,
|
||||
},
|
||||
)
|
||||
tool_card.artifacts.append(
|
||||
{
|
||||
"artifact_type": ArtifactType.WEB_PAGES.value,
|
||||
"artifact_id": artifact_id,
|
||||
}
|
||||
)
|
||||
|
||||
return f"""\
|
||||
\n\n**🔎 Search Results**\n\n
|
||||
```tool_card
|
||||
{json.dumps(tool_card.model_dump(), ensure_ascii=False, indent=2)}
|
||||
```\n
|
||||
"""
|
||||
|
||||
|
||||
class ToolResultParserFactory:
|
||||
def get_parser(self, tool_type: str, tool_name: str):
|
||||
if "search" in tool_name and ("search" in tool_name or tool_name == None):
|
||||
return SearchToolResultParser()
|
||||
else:
|
||||
return DefaultToolResultParser()
|
||||
|
||||
|
||||
@dataclass
|
||||
class AWorldWebAgentUI(AworldUI):
|
||||
session_id: str = Field(default="", description="session id")
|
||||
workspace: WorkSpace = Field(default=None, description="workspace")
|
||||
tool_result_parser_factory: ToolResultParserFactory = Field(
|
||||
default=ToolResultParserFactory, description="tool result parser factory"
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
session_id: str = None,
|
||||
workspace: WorkSpace = None,
|
||||
tool_result_parser_factory: ToolResultParserFactory = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Initialize MarkdownAworldUI
|
||||
Args:"""
|
||||
super().__init__(**kwargs)
|
||||
self.session_id = session_id
|
||||
self.workspace = workspace
|
||||
self.tool_result_parser_factory = (
|
||||
tool_result_parser_factory or ToolResultParserFactory()
|
||||
)
|
||||
|
||||
@override
|
||||
async def message_output(self, __output__: MessageOutput):
|
||||
"""
|
||||
Returns an async generator that yields each message item.
|
||||
"""
|
||||
# Sentinel object for queue completion
|
||||
_SENTINEL = object()
|
||||
|
||||
async def async_generator():
|
||||
async def __log_item(item):
|
||||
await queue.put(item)
|
||||
|
||||
from asyncio import Queue
|
||||
|
||||
queue = Queue()
|
||||
|
||||
async def consume_all():
|
||||
# Consume all relevant generators
|
||||
if __output__.reason_generator or __output__.response_generator:
|
||||
if __output__.reason_generator:
|
||||
await consume_content(__output__.reason_generator, __log_item)
|
||||
if __output__.response_generator:
|
||||
await consume_content(__output__.response_generator, __log_item)
|
||||
else:
|
||||
await consume_content(__output__.reasoning, __log_item)
|
||||
await consume_content(__output__.response, __log_item)
|
||||
# Only after all are done, put the sentinel
|
||||
await queue.put(_SENTINEL)
|
||||
|
||||
# Start the consumer in the background
|
||||
import asyncio
|
||||
|
||||
consumer_task = asyncio.create_task(consume_all())
|
||||
|
||||
while True:
|
||||
item = await queue.get()
|
||||
if item is _SENTINEL:
|
||||
break
|
||||
yield item
|
||||
await consumer_task # Ensure background task is finished
|
||||
|
||||
return async_generator()
|
||||
|
||||
@override
|
||||
async def tool_result(self, output: ToolResultOutput):
|
||||
"""
|
||||
tool_result
|
||||
"""
|
||||
parser = self.tool_result_parser_factory.get_parser(
|
||||
output.tool_type, output.tool_name
|
||||
)
|
||||
return await parser.parse(output, workspace=self.workspace)
|
||||
|
||||
@override
|
||||
async def step(self, output: StepOutput):
|
||||
emptyLine = "\n\n"
|
||||
if output.status == "START":
|
||||
return f"\n\n # {output.show_name} \n\n"
|
||||
elif output.status == "FINISHED":
|
||||
return f"{emptyLine}"
|
||||
elif output.status == "FAILED":
|
||||
return f"\n\n{output.name} 💥FAILED: reason is {output.data} {emptyLine}"
|
||||
else:
|
||||
return f"\n\n{output.name} ❓❓❓UNKNOWN#{output.status} {emptyLine}"
|
||||
|
||||
@override
|
||||
async def custom_output(self, output: Output):
|
||||
return output.data
|
||||
@@ -0,0 +1,175 @@
|
||||
import os
|
||||
import logging
|
||||
import traceback
|
||||
import asyncio
|
||||
import re
|
||||
import json
|
||||
import pickle
|
||||
from asyncio.tasks import Task
|
||||
from aworld.config.conf import AgentConfig
|
||||
from aworld.agents.llm_agent import Agent
|
||||
from typing import Dict, Union
|
||||
|
||||
from aworld.core.context.base import Context
|
||||
from aworld.utils.run_util import exec_agent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SimpleSummaryCache:
|
||||
def __init__(self) -> None:
|
||||
self._cache_file = os.path.join(os.curdir, "data", "trace_summary_cache.pkl")
|
||||
self._cache: Dict[str, str] = {}
|
||||
self._load_cache()
|
||||
|
||||
def _load_cache(self):
|
||||
if os.path.exists(self._cache_file):
|
||||
try:
|
||||
with open(self._cache_file, "rb") as f:
|
||||
self._cache = pickle.load(f)
|
||||
except (pickle.PickleError, EOFError):
|
||||
logger.warning("Cache file is corrupted, creating new cache")
|
||||
if self._cache_file.exists():
|
||||
self._cache_file.unlink()
|
||||
|
||||
def _save_cache(self):
|
||||
serializable_cache = {
|
||||
k: v for k, v in self._cache.items() if not isinstance(v, Task)
|
||||
}
|
||||
try:
|
||||
with open(self._cache_file, "wb") as f:
|
||||
pickle.dump(serializable_cache, f)
|
||||
except pickle.PickleError:
|
||||
logger.error("Failed to save cache")
|
||||
|
||||
def add_to_cache(self, trace_id: str, value: Union[str, Task]):
|
||||
self._cache[trace_id] = value
|
||||
if not isinstance(value, Task):
|
||||
self._save_cache()
|
||||
|
||||
def get_value(self, trace_id: str) -> Union[str, Task]:
|
||||
return self._cache.get(trace_id)
|
||||
|
||||
def trace_exists(self, trace_id: str) -> bool:
|
||||
return trace_id in self._cache
|
||||
|
||||
|
||||
# _trace_summary_cache: Dict[str, Union[str, Task]] = {}
|
||||
_trace_summary_cache = SimpleSummaryCache()
|
||||
|
||||
trace_sys_prompt = "You are a helpful tracking summary agent."
|
||||
|
||||
trace_prompt = """
|
||||
you can use tracking tools to obtain tracking data and then summarize the main tasks completed by each agent and their token usage.
|
||||
You can identify which spans are agents, which spans are tool calls, and which spans are large model calls based on the following criteria:
|
||||
1 Agent span: the prefix for 'name' is 'event.agent.'
|
||||
2 LLM span: The prefix for 'name' is 'llm.'
|
||||
3 Tool span: The prefix for 'name' is 'event.tool.'
|
||||
|
||||
requirement:
|
||||
1. Please summarize and output separately for agents with different event.id.
|
||||
2. Agent Span with the same name but different event.id are also considered as different agents.
|
||||
3. There may be a parent-child relationship between agents. Please select the LLM span and Tool span from the nearest child span to the current agent for summarizing.
|
||||
4. Ensure that all agent spans have their own independent summaries, and the number of summaries is exactly the same as the number of agent spans. For example: {{"name":"event.agent.a","attributes":{{"event.id":"111"}},"children":[{{"name":"llm.gpt-4o"}},{{"name":"event.tool.1","children":[{{"name":"event.agent.a","attributes":{{"event.id":"222"}},"children":[{{"name":"llm.gpt-4o"}}]}}]}}]}}, both of the above two agent names are event.agent.a, but event.id is different and needs to be summarized separately for 111 and 222. So you need to identify all agent spans without any omissions, which is very important.
|
||||
5. Please output in the following standard JSON format without any additional explanatory text:
|
||||
[{{"agent":"947cc4c1b7ed406ab7fbf38b9d2b1f5a",,"summary":"xxx"}},{{}}]
|
||||
6. Pay attention to controlling the length of the summary, so that the overall output does not exceed your output length limit.
|
||||
Here are the trace_id: {task}
|
||||
"""
|
||||
|
||||
agent_config = None
|
||||
|
||||
|
||||
async def _do_summarize_trace(trace_id: str):
|
||||
logger.info(f"_do_summarize_trace trace_id: {trace_id}")
|
||||
global agent_config
|
||||
trace_agent = Agent(
|
||||
conf=agent_config,
|
||||
name="trace_agent",
|
||||
system_prompt=trace_sys_prompt,
|
||||
agent_prompt=trace_prompt,
|
||||
tool_names=["trace"],
|
||||
feedback_tool_result=True,
|
||||
)
|
||||
|
||||
if trace_agent.conf.llm_config.llm_api_key is None:
|
||||
logger.warning(
|
||||
"LLM_API_KEY_TRACE is not set, trace summarize will not be executed."
|
||||
)
|
||||
return ""
|
||||
try:
|
||||
res = await exec_agent(trace_id, trace_agent, Context())
|
||||
summary = _fetch_json_from_result(res.answer)
|
||||
_trace_summary_cache.add_to_cache(trace_id, summary)
|
||||
return summary
|
||||
except Exception as e:
|
||||
logger.error(traceback.format_exc())
|
||||
|
||||
|
||||
def summarize_trace(trace_id: str):
|
||||
global agent_config
|
||||
if agent_config is None:
|
||||
llm_provider = os.getenv("LLM_PROVIDER_TRACE", "openai")
|
||||
llm_model_name = os.getenv("LLM_MODEL_NAME_TRACE", None)
|
||||
llm_base_url = os.getenv("LLM_BASE_URL_TRACE", None)
|
||||
llm_api_key = os.getenv("LLM_API_KEY_TRACE", None)
|
||||
|
||||
if (
|
||||
not llm_provider
|
||||
or not llm_model_name
|
||||
or not llm_base_url
|
||||
or not llm_api_key
|
||||
):
|
||||
logger.warning(
|
||||
"LLM_MODEL_NAME_TRACE, LLM_BASE_URL_TRACE, LLM_API_KEY_TRACE is not set, trace summarize will not be executed."
|
||||
)
|
||||
return
|
||||
|
||||
agent_config = AgentConfig(
|
||||
llm_provider=os.getenv("LLM_PROVIDER_TRACE", "openai"),
|
||||
llm_model_name=os.getenv("LLM_MODEL_NAME_TRACE", None),
|
||||
llm_base_url=os.getenv("LLM_BASE_URL_TRACE", None),
|
||||
llm_api_key=os.getenv("LLM_API_KEY_TRACE", None),
|
||||
)
|
||||
llm_config = agent_config.llm_config
|
||||
if not _trace_summary_cache.trace_exists(trace_id):
|
||||
if (
|
||||
llm_config.llm_api_key is None
|
||||
or not llm_config.llm_base_url
|
||||
or not llm_config.llm_model_name
|
||||
):
|
||||
logger.warning(
|
||||
"LLM_MODEL_NAME_TRACE, LLM_BASE_URL_TRACE, LLM_API_KEY_TRACE is not set, trace summarize will not be executed."
|
||||
)
|
||||
return
|
||||
|
||||
task = asyncio.create_task(_do_summarize_trace(trace_id))
|
||||
_trace_summary_cache.add_to_cache(trace_id, task)
|
||||
|
||||
|
||||
async def get_summarize_trace(trace_id: str):
|
||||
if not _trace_summary_cache.trace_exists(trace_id):
|
||||
return None
|
||||
cached_value = _trace_summary_cache.get_value(trace_id)
|
||||
if isinstance(cached_value, Task):
|
||||
# try:
|
||||
# result = await cached_value
|
||||
# if isinstance(result, Task):
|
||||
# result = await result
|
||||
# _trace_summary_cache[trace_id] = _fetch_json_from_result(result)
|
||||
# except Exception as e:
|
||||
# logger.error(traceback.format_exc())
|
||||
return None
|
||||
return cached_value
|
||||
|
||||
|
||||
def _fetch_json_from_result(input_str):
|
||||
json_match = re.search(r"\[.*\]", input_str, re.DOTALL)
|
||||
if json_match:
|
||||
json_str = json_match.group(0)
|
||||
try:
|
||||
json.loads(json_str)
|
||||
return json_str
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f"_fetch_json_from_result json_str: {json_str} error: {e}")
|
||||
return ""
|
||||
@@ -0,0 +1,26 @@
|
||||
import subprocess
|
||||
import logging
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def build_webui(force_rebuild: bool = False) -> str:
|
||||
webui_path = Path(__file__).parent.parent / "web" / "webui"
|
||||
static_path = webui_path / "dist"
|
||||
|
||||
if (not static_path.exists()) or force_rebuild:
|
||||
logger.warning(f"Build WebUI at {webui_path}")
|
||||
|
||||
try:
|
||||
subprocess.check_call(
|
||||
["sh", "-c", "npm install && npm run build"],
|
||||
cwd=webui_path,
|
||||
)
|
||||
logger.info("WebUI build successfully")
|
||||
except:
|
||||
logger.error(f"Failed to build WebUI at {webui_path}")
|
||||
sys.exit(1)
|
||||
|
||||
return static_path
|
||||
Reference in New Issue
Block a user