add preliminary support for llm agents (#432)

This commit is contained in:
Adil Hafeez 2025-03-19 15:21:34 -07:00 committed by GitHub
parent 8d66fefded
commit 84cd1df7bf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
29 changed files with 1388 additions and 121 deletions

View file

@ -3,6 +3,8 @@ from openai import OpenAI
from src.commons.utils import get_model_server_logger
from src.core.guardrails import get_guardrail_handler
from src.core.function_calling import (
ArchAgentConfig,
ArchAgentHandler,
ArchIntentConfig,
ArchIntentHandler,
ArchFunctionConfig,
@ -18,10 +20,12 @@ logger = get_model_server_logger()
ARCH_ENDPOINT = os.getenv("ARCH_ENDPOINT", "https://archfc.katanemo.dev/v1")
ARCH_API_KEY = "EMPTY"
ARCH_CLIENT = OpenAI(base_url=ARCH_ENDPOINT, api_key=ARCH_API_KEY)
ARCH_AGENT_CLIENT = ARCH_CLIENT
# Define model names
ARCH_INTENT_MODEL_ALIAS = "Arch-Intent"
ARCH_FUNCTION_MODEL_ALIAS = "Arch-Function"
ARCH_AGENT_MODEL_ALIAS = ARCH_FUNCTION_MODEL_ALIAS
ARCH_GUARD_MODEL_ALIAS = "katanemo/Arch-Guard"
# Define model handlers
@ -32,5 +36,8 @@ handler_map = {
"Arch-Function": ArchFunctionHandler(
ARCH_CLIENT, ARCH_FUNCTION_MODEL_ALIAS, ArchFunctionConfig
),
"Arch-Agent": ArchAgentHandler(
ARCH_AGENT_CLIENT, ARCH_AGENT_MODEL_ALIAS, ArchAgentConfig
),
"Arch-Guard": get_guardrail_handler(ARCH_GUARD_MODEL_ALIAS),
}