146 lines
6.0 KiB
Python
146 lines
6.0 KiB
Python
import uuid
|
|
from functools import partial
|
|
from typing import List
|
|
|
|
from fastapi import APIRouter, Body, Depends, HTTPException
|
|
from pydantic import BaseModel, Field
|
|
|
|
from memgpt.constants import BASE_TOOLS
|
|
from memgpt.memory import ChatMemory
|
|
from memgpt.models.pydantic_models import (
|
|
AgentStateModel,
|
|
EmbeddingConfigModel,
|
|
LLMConfigModel,
|
|
PresetModel,
|
|
)
|
|
from memgpt.server.rest_api.auth_token import get_current_user
|
|
from memgpt.server.rest_api.interface import QueuingInterface
|
|
from memgpt.server.server import SyncServer
|
|
from memgpt.settings import settings
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
class ListAgentsResponse(BaseModel):
|
|
num_agents: int = Field(..., description="The number of agents available to the user.")
|
|
# TODO make return type List[AgentStateModel]
|
|
# also return - presets: List[PresetModel]
|
|
agents: List[dict] = Field(..., description="List of agent configurations.")
|
|
|
|
|
|
class CreateAgentRequest(BaseModel):
|
|
# TODO: modify this (along with front end)
|
|
config: dict = Field(..., description="The agent configuration object.")
|
|
|
|
|
|
class CreateAgentResponse(BaseModel):
|
|
agent_state: AgentStateModel = Field(..., description="The state of the newly created agent.")
|
|
preset: PresetModel = Field(..., description="The preset that the agent was created from.")
|
|
|
|
|
|
def setup_agents_index_router(server: SyncServer, interface: QueuingInterface, password: str):
|
|
get_current_user_with_server = partial(partial(get_current_user, server), password)
|
|
|
|
@router.get("/agents", tags=["agents"], response_model=ListAgentsResponse)
|
|
def list_agents(
|
|
user_id: uuid.UUID = Depends(get_current_user_with_server),
|
|
):
|
|
"""
|
|
List all agents associated with a given user.
|
|
|
|
This endpoint retrieves a list of all agents and their configurations associated with the specified user ID.
|
|
"""
|
|
interface.clear()
|
|
agents_data = server.list_agents_legacy(user_id=user_id)
|
|
return ListAgentsResponse(**agents_data)
|
|
|
|
@router.post("/agents", tags=["agents"], response_model=CreateAgentResponse)
|
|
def create_agent(
|
|
request: CreateAgentRequest = Body(...),
|
|
user_id: uuid.UUID = Depends(get_current_user_with_server),
|
|
):
|
|
"""
|
|
Create a new agent with the specified configuration.
|
|
"""
|
|
interface.clear()
|
|
|
|
# Parse request
|
|
# TODO: don't just use JSON in the future
|
|
human_name = request.config["human_name"] if "human_name" in request.config else None
|
|
human = request.config["human"] if "human" in request.config else None
|
|
persona_name = request.config["persona_name"] if "persona_name" in request.config else None
|
|
persona = request.config["persona"] if "persona" in request.config else None
|
|
request.config["preset"] if ("preset" in request.config and request.config["preset"]) else settings.default_preset
|
|
tool_names = request.config["function_names"] if ("function_names" in request.config and request.config["function_names"]) else None
|
|
metadata = request.config["metadata"] if "metadata" in request.config else {}
|
|
metadata["human"] = human_name
|
|
metadata["persona"] = persona_name
|
|
|
|
# TODO: remove this -- should be added based on create agent fields
|
|
if isinstance(tool_names, str): # TODO: fix this on clinet side?
|
|
tool_names = tool_names.split(",")
|
|
if tool_names is None or tool_names == "":
|
|
tool_names = []
|
|
for name in BASE_TOOLS: # TODO: remove this
|
|
if name not in tool_names:
|
|
tool_names.append(name)
|
|
assert isinstance(tool_names, list), "Tool names must be a list of strings."
|
|
|
|
# TODO: eventually remove this - should support general memory at the REST endpoint
|
|
# TODO: the REST server should add default memory tools at startup time
|
|
memory = ChatMemory(persona=persona, human=human)
|
|
|
|
try:
|
|
agent_state = server.create_agent(
|
|
user_id=user_id,
|
|
# **request.config
|
|
# TODO turn into a pydantic model
|
|
name=request.config["name"],
|
|
memory=memory,
|
|
system=request.config.get("system", None),
|
|
# persona_name=persona_name,
|
|
# human_name=human_name,
|
|
# persona=persona,
|
|
# human=human,
|
|
# llm_config=LLMConfigModel(
|
|
# model=request.config['model'],
|
|
# )
|
|
# tools
|
|
tools=tool_names,
|
|
metadata=metadata,
|
|
# function_names=request.config["function_names"].split(",") if "function_names" in request.config else None,
|
|
)
|
|
llm_config = LLMConfigModel(**vars(agent_state.llm_config))
|
|
embedding_config = EmbeddingConfigModel(**vars(agent_state.embedding_config))
|
|
|
|
return CreateAgentResponse(
|
|
agent_state=AgentStateModel(
|
|
id=agent_state.id,
|
|
name=agent_state.name,
|
|
user_id=agent_state.user_id,
|
|
llm_config=llm_config,
|
|
embedding_config=embedding_config,
|
|
state=agent_state.state,
|
|
created_at=int(agent_state.created_at.timestamp()),
|
|
tools=agent_state.tools,
|
|
system=agent_state.system,
|
|
metadata=agent_state._metadata,
|
|
),
|
|
preset=PresetModel( # TODO: remove (placeholder to avoid breaking frontend)
|
|
name="dummy_preset",
|
|
id=agent_state.id,
|
|
user_id=agent_state.user_id,
|
|
description="",
|
|
created_at=agent_state.created_at,
|
|
system=agent_state.system,
|
|
persona="",
|
|
human="",
|
|
functions_schema=[],
|
|
),
|
|
)
|
|
except Exception as e:
|
|
print(str(e))
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
return router
|