flesh out memory banks API

This commit is contained in:
Ashwin Bharambe 2024-08-23 06:38:15 -07:00
parent 31289e3f47
commit 77d6055d9f
11 changed files with 1792 additions and 974 deletions

View file

@ -8,7 +8,7 @@
import copy
import uuid
from datetime import datetime
from typing import AsyncGenerator, List, Optional
from typing import AsyncGenerator, List
from llama_models.llama3.api.datatypes import ToolPromptFormat
@ -326,7 +326,7 @@ class ChatAgent(ShieldRunnerMixin):
req = ChatCompletionRequest(
model=self.agent_config.model,
messages=input_messages,
tools=self.agent_config.available_tools,
tools=self.agent_config.tools,
tool_prompt_format=self.agent_config.tool_prompt_format,
stream=True,
sampling_params=sampling_params,

View file

@ -24,7 +24,7 @@ from llama_toolchain.tools.builtin import (
)
from llama_toolchain.tools.safety import with_safety
from .agent_instance import AgentInstance, ChatAgent
from .agent_instance import ChatAgent
from .config import MetaReferenceImplConfig
@ -71,11 +71,11 @@ class MetaReferenceAgenticSystemImpl(AgenticSystem):
self,
request: AgenticSystemCreateRequest,
) -> AgenticSystemCreateResponse:
system_id = str(uuid.uuid4())
agent_id = str(uuid.uuid4())
builtin_tools = []
cfg = request.agent_config
for dfn in cfg.available_tools:
for dfn in cfg.tools:
if isinstance(dfn.tool_name, BuiltinTool):
if dfn.tool_name == BuiltinTool.wolfram_alpha:
key = self.config.wolfram_api_key
@ -102,7 +102,7 @@ class MetaReferenceAgenticSystemImpl(AgenticSystem):
)
)
AGENT_INSTANCES_BY_ID[system_id] = ChatAgent(
AGENT_INSTANCES_BY_ID[agent_id] = ChatAgent(
agent_config=cfg,
inference_api=self.inference_api,
safety_api=self.safety_api,
@ -111,16 +111,16 @@ class MetaReferenceAgenticSystemImpl(AgenticSystem):
)
return AgenticSystemCreateResponse(
system_id=system_id,
agent_id=agent_id,
)
async def create_agentic_system_session(
self,
request: AgenticSystemSessionCreateRequest,
) -> AgenticSystemSessionCreateResponse:
system_id = request.system_id
assert system_id in AGENT_INSTANCES_BY_ID, f"System {system_id} not found"
agent = AGENT_INSTANCES_BY_ID[system_id]
agent_id = request.agent_id
assert agent_id in AGENT_INSTANCES_BY_ID, f"System {agent_id} not found"
agent = AGENT_INSTANCES_BY_ID[agent_id]
session = agent.create_session(request.session_name)
return AgenticSystemSessionCreateResponse(
@ -131,9 +131,9 @@ class MetaReferenceAgenticSystemImpl(AgenticSystem):
self,
request: AgenticSystemTurnCreateRequest,
) -> AsyncGenerator:
system_id = request.system_id
assert system_id in AGENT_INSTANCES_BY_ID, f"System {system_id} not found"
agent = AGENT_INSTANCES_BY_ID[system_id]
agent_id = request.agent_id
assert agent_id in AGENT_INSTANCES_BY_ID, f"System {agent_id} not found"
agent = AGENT_INSTANCES_BY_ID[agent_id]
assert (
request.session_id in agent.sessions

View file

@ -19,7 +19,7 @@ from llama_toolchain.inference.api import Message
async def execute_with_custom_tools(
system: AgenticSystem,
system_id: str,
agent_id: str,
session_id: str,
messages: List[Message],
custom_tools: List[Any],
@ -35,7 +35,7 @@ async def execute_with_custom_tools(
n_iter += 1
request = AgenticSystemTurnCreateRequest(
system_id=system_id,
agent_id=agent_id,
session_id=session_id,
messages=current_messages,
stream=stream,