mirror of
https://github.com/arc53/DocsGPT.git
synced 2026-05-07 06:30:03 +00:00
1073 lines
42 KiB
Python
1073 lines
42 KiB
Python
import datetime
|
|
import json
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional, Set
|
|
|
|
from bson.dbref import DBRef
|
|
|
|
from bson.objectid import ObjectId
|
|
|
|
from application.agents.agent_creator import AgentCreator
|
|
from application.api.answer.services.compression import CompressionOrchestrator
|
|
from application.api.answer.services.compression.token_counter import TokenCounter
|
|
from application.api.answer.services.conversation_service import ConversationService
|
|
from application.api.answer.services.prompt_renderer import PromptRenderer
|
|
from application.core.model_utils import (
|
|
get_api_key_for_provider,
|
|
get_default_model_id,
|
|
get_provider_from_model_id,
|
|
validate_model_id,
|
|
)
|
|
from application.core.mongo_db import MongoDB
|
|
from application.core.settings import settings
|
|
from application.retriever.retriever_creator import RetrieverCreator
|
|
from application.utils import (
|
|
calculate_doc_token_budget,
|
|
limit_chat_history,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def get_prompt(prompt_id: str, prompts_collection=None) -> str:
|
|
"""
|
|
Get a prompt by preset name or MongoDB ID
|
|
"""
|
|
current_dir = Path(__file__).resolve().parents[3]
|
|
prompts_dir = current_dir / "prompts"
|
|
|
|
# Maps for classic agent types
|
|
CLASSIC_PRESETS = {
|
|
"default": "chat_combine_default.txt",
|
|
"creative": "chat_combine_creative.txt",
|
|
"strict": "chat_combine_strict.txt",
|
|
"reduce": "chat_reduce_prompt.txt",
|
|
}
|
|
|
|
# Agentic counterparts — same styles, but with search tool instructions
|
|
AGENTIC_PRESETS = {
|
|
"default": "agentic/default.txt",
|
|
"creative": "agentic/creative.txt",
|
|
"strict": "agentic/strict.txt",
|
|
}
|
|
|
|
preset_mapping = {**CLASSIC_PRESETS, **{f"agentic_{k}": v for k, v in AGENTIC_PRESETS.items()}}
|
|
|
|
if prompt_id in preset_mapping:
|
|
file_path = os.path.join(prompts_dir, preset_mapping[prompt_id])
|
|
try:
|
|
with open(file_path, "r") as f:
|
|
return f.read()
|
|
except FileNotFoundError:
|
|
raise FileNotFoundError(f"Prompt file not found: {file_path}")
|
|
try:
|
|
if prompts_collection is None:
|
|
mongo = MongoDB.get_client()
|
|
db = mongo[settings.MONGO_DB_NAME]
|
|
prompts_collection = db["prompts"]
|
|
prompt_doc = prompts_collection.find_one({"_id": ObjectId(prompt_id)})
|
|
if not prompt_doc:
|
|
raise ValueError(f"Prompt with ID {prompt_id} not found")
|
|
return prompt_doc["content"]
|
|
except Exception as e:
|
|
raise ValueError(f"Invalid prompt ID: {prompt_id}") from e
|
|
|
|
|
|
class StreamProcessor:
|
|
def __init__(
|
|
self, request_data: Dict[str, Any], decoded_token: Optional[Dict[str, Any]]
|
|
):
|
|
mongo = MongoDB.get_client()
|
|
self.db = mongo[settings.MONGO_DB_NAME]
|
|
self.agents_collection = self.db["agents"]
|
|
self.attachments_collection = self.db["attachments"]
|
|
self.prompts_collection = self.db["prompts"]
|
|
|
|
self.data = request_data
|
|
self.decoded_token = decoded_token
|
|
self.initial_user_id = (
|
|
self.decoded_token.get("sub") if self.decoded_token is not None else None
|
|
)
|
|
self.conversation_id = self.data.get("conversation_id")
|
|
self.source = {}
|
|
self.all_sources = []
|
|
self.attachments = []
|
|
self.history = []
|
|
self.retrieved_docs = []
|
|
self.agent_config = {}
|
|
self.retriever_config = {}
|
|
self.is_shared_usage = False
|
|
self.shared_token = None
|
|
self.agent_id = self.data.get("agent_id")
|
|
self.agent_key = None
|
|
self.model_id: Optional[str] = None
|
|
self.conversation_service = ConversationService()
|
|
self.compression_orchestrator = CompressionOrchestrator(
|
|
self.conversation_service
|
|
)
|
|
self.prompt_renderer = PromptRenderer()
|
|
self._prompt_content: Optional[str] = None
|
|
self._required_tool_actions: Optional[Dict[str, Set[Optional[str]]]] = None
|
|
self.compressed_summary: Optional[str] = None
|
|
self.compressed_summary_tokens: int = 0
|
|
self._agent_data: Optional[Dict[str, Any]] = None
|
|
|
|
def initialize(self):
|
|
"""Initialize all required components for processing"""
|
|
self._configure_agent()
|
|
self._validate_and_set_model()
|
|
self._configure_source()
|
|
self._configure_retriever()
|
|
self._load_conversation_history()
|
|
self._process_attachments()
|
|
|
|
def build_agent(self, question: str):
|
|
"""One call to go from request data to a ready-to-run agent.
|
|
|
|
Combines initialize(), pre_fetch_docs(), pre_fetch_tools(), and
|
|
create_agent() into a single convenience method.
|
|
"""
|
|
self.initialize()
|
|
|
|
agent_type = self.agent_config.get("agent_type", "classic")
|
|
|
|
# Agentic/research agents skip pre-fetch — the LLM searches on-demand via tools
|
|
if agent_type in ("agentic", "research"):
|
|
tools_data = self.pre_fetch_tools()
|
|
return self.create_agent(tools_data=tools_data)
|
|
|
|
docs_together, docs_list = self.pre_fetch_docs(question)
|
|
tools_data = self.pre_fetch_tools()
|
|
return self.create_agent(
|
|
docs_together=docs_together,
|
|
docs=docs_list,
|
|
tools_data=tools_data,
|
|
)
|
|
|
|
def _load_conversation_history(self):
|
|
"""Load conversation history either from DB or request"""
|
|
if self.conversation_id and self.initial_user_id:
|
|
conversation = self.conversation_service.get_conversation(
|
|
self.conversation_id, self.initial_user_id
|
|
)
|
|
if not conversation:
|
|
raise ValueError("Conversation not found or unauthorized")
|
|
|
|
# Check if compression is enabled and needed
|
|
if settings.ENABLE_CONVERSATION_COMPRESSION:
|
|
self._handle_compression(conversation)
|
|
else:
|
|
# Original behavior - load all history (include metadata if present)
|
|
self.history = [
|
|
{
|
|
"prompt": query["prompt"],
|
|
"response": query["response"],
|
|
**(
|
|
{"metadata": query["metadata"]}
|
|
if "metadata" in query
|
|
else {}
|
|
),
|
|
}
|
|
for query in conversation.get("queries", [])
|
|
]
|
|
else:
|
|
self.history = limit_chat_history(
|
|
json.loads(self.data.get("history", "[]")), model_id=self.model_id
|
|
)
|
|
|
|
def _handle_compression(self, conversation: Dict[str, Any]):
|
|
"""Handle conversation compression logic using orchestrator."""
|
|
try:
|
|
result = self.compression_orchestrator.compress_if_needed(
|
|
conversation_id=self.conversation_id,
|
|
user_id=self.initial_user_id,
|
|
model_id=self.model_id,
|
|
decoded_token=self.decoded_token,
|
|
)
|
|
|
|
if not result.success:
|
|
logger.error(f"Compression failed: {result.error}, using full history")
|
|
self.history = [
|
|
{
|
|
"prompt": query["prompt"],
|
|
"response": query["response"],
|
|
**({"metadata": query["metadata"]} if "metadata" in query else {}),
|
|
}
|
|
for query in conversation.get("queries", [])
|
|
]
|
|
return
|
|
|
|
if result.compression_performed and result.compressed_summary:
|
|
self.compressed_summary = result.compressed_summary
|
|
self.compressed_summary_tokens = TokenCounter.count_message_tokens(
|
|
[{"content": result.compressed_summary}]
|
|
)
|
|
logger.info(
|
|
f"Using compressed summary ({self.compressed_summary_tokens} tokens) "
|
|
f"+ {len(result.recent_queries)} recent messages"
|
|
)
|
|
|
|
self.history = result.as_history()
|
|
# Preserve metadata from recent queries (as_history only has prompt/response)
|
|
recent = result.recent_queries if result.recent_queries else conversation.get("queries", [])
|
|
for i, entry in enumerate(self.history):
|
|
# Match by index from the end of recent queries
|
|
offset = len(recent) - len(self.history)
|
|
qi = offset + i
|
|
if 0 <= qi < len(recent) and "metadata" in recent[qi]:
|
|
entry["metadata"] = recent[qi]["metadata"]
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Error handling compression, falling back to standard history: {str(e)}",
|
|
exc_info=True,
|
|
)
|
|
self.history = [
|
|
{
|
|
"prompt": query["prompt"],
|
|
"response": query["response"],
|
|
**({"metadata": query["metadata"]} if "metadata" in query else {}),
|
|
}
|
|
for query in conversation.get("queries", [])
|
|
]
|
|
|
|
def _process_attachments(self):
|
|
"""Process any attachments in the request"""
|
|
attachment_ids = self.data.get("attachments", [])
|
|
self.attachments = self._get_attachments_content(
|
|
attachment_ids, self.initial_user_id
|
|
)
|
|
|
|
def _get_attachments_content(self, attachment_ids, user_id):
|
|
if not attachment_ids:
|
|
return []
|
|
attachments = []
|
|
for attachment_id in attachment_ids:
|
|
try:
|
|
attachment_doc = self.attachments_collection.find_one(
|
|
{"_id": ObjectId(attachment_id), "user": user_id}
|
|
)
|
|
if attachment_doc:
|
|
attachments.append(attachment_doc)
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Error retrieving attachment {attachment_id}: {e}", exc_info=True
|
|
)
|
|
return attachments
|
|
|
|
def _validate_and_set_model(self):
|
|
"""Validate and set model_id from request"""
|
|
from application.core.model_settings import ModelRegistry
|
|
|
|
requested_model = self.data.get("model_id")
|
|
|
|
if requested_model:
|
|
if not validate_model_id(requested_model):
|
|
registry = ModelRegistry.get_instance()
|
|
available_models = [m.id for m in registry.get_enabled_models()]
|
|
raise ValueError(
|
|
f"Invalid model_id '{requested_model}'. "
|
|
f"Available models: {', '.join(available_models[:5])}"
|
|
+ (
|
|
f" and {len(available_models) - 5} more"
|
|
if len(available_models) > 5
|
|
else ""
|
|
)
|
|
)
|
|
self.model_id = requested_model
|
|
else:
|
|
agent_default_model = self.agent_config.get("default_model_id", "")
|
|
if agent_default_model and validate_model_id(agent_default_model):
|
|
self.model_id = agent_default_model
|
|
else:
|
|
self.model_id = get_default_model_id()
|
|
|
|
def _get_agent_key(self, agent_id: Optional[str], user_id: Optional[str]) -> tuple:
|
|
"""Get API key for agent with access control"""
|
|
if not agent_id:
|
|
return None, False, None
|
|
try:
|
|
agent = self.agents_collection.find_one({"_id": ObjectId(agent_id)})
|
|
if agent is None:
|
|
raise Exception("Agent not found")
|
|
is_owner = agent.get("user") == user_id
|
|
is_shared_with_user = agent.get(
|
|
"shared_publicly", False
|
|
) or user_id in agent.get("shared_with", [])
|
|
|
|
if not (is_owner or is_shared_with_user):
|
|
raise Exception("Unauthorized access to the agent")
|
|
if is_owner:
|
|
self.agents_collection.update_one(
|
|
{"_id": ObjectId(agent_id)},
|
|
{
|
|
"$set": {
|
|
"lastUsedAt": datetime.datetime.now(datetime.timezone.utc)
|
|
}
|
|
},
|
|
)
|
|
return str(agent["key"]), not is_owner, agent.get("shared_token")
|
|
except Exception as e:
|
|
logger.error(f"Error in get_agent_key: {str(e)}", exc_info=True)
|
|
raise
|
|
|
|
def _get_data_from_api_key(self, api_key: str) -> Dict[str, Any]:
|
|
data = self.agents_collection.find_one({"key": api_key})
|
|
if not data:
|
|
raise Exception("Invalid API Key, please generate a new key", 401)
|
|
source = data.get("source")
|
|
if isinstance(source, DBRef):
|
|
source_doc = self.db.dereference(source)
|
|
if source_doc:
|
|
data["source"] = str(source_doc["_id"])
|
|
data["retriever"] = source_doc.get("retriever", data.get("retriever"))
|
|
data["chunks"] = source_doc.get("chunks", data.get("chunks"))
|
|
else:
|
|
data["source"] = None
|
|
elif source == "default":
|
|
data["source"] = "default"
|
|
else:
|
|
data["source"] = None
|
|
|
|
sources = data.get("sources", [])
|
|
if sources and isinstance(sources, list):
|
|
sources_list = []
|
|
for i, source_ref in enumerate(sources):
|
|
if source_ref == "default":
|
|
processed_source = {
|
|
"id": "default",
|
|
"retriever": "classic",
|
|
"chunks": data.get("chunks", "2"),
|
|
}
|
|
sources_list.append(processed_source)
|
|
elif isinstance(source_ref, DBRef):
|
|
source_doc = self.db.dereference(source_ref)
|
|
if source_doc:
|
|
processed_source = {
|
|
"id": str(source_doc["_id"]),
|
|
"retriever": source_doc.get("retriever", "classic"),
|
|
"chunks": source_doc.get("chunks", data.get("chunks", "2")),
|
|
}
|
|
sources_list.append(processed_source)
|
|
data["sources"] = sources_list
|
|
else:
|
|
data["sources"] = []
|
|
|
|
data["default_model_id"] = data.get("default_model_id", "")
|
|
|
|
return data
|
|
|
|
def _configure_source(self):
|
|
"""Configure the source based on agent data.
|
|
|
|
The literal string ``"default"`` is a placeholder meaning "no
|
|
ingested source" and is normalized to an empty source so that no
|
|
retrieval is attempted.
|
|
"""
|
|
if self._agent_data:
|
|
agent_data = self._agent_data
|
|
|
|
if agent_data.get("sources") and len(agent_data["sources"]) > 0:
|
|
source_ids = [
|
|
source["id"]
|
|
for source in agent_data["sources"]
|
|
if source.get("id") and source["id"] != "default"
|
|
]
|
|
if source_ids:
|
|
self.source = {"active_docs": source_ids}
|
|
else:
|
|
self.source = {}
|
|
self.all_sources = [
|
|
s for s in agent_data["sources"] if s.get("id") != "default"
|
|
]
|
|
elif agent_data.get("source") and agent_data["source"] != "default":
|
|
self.source = {"active_docs": agent_data["source"]}
|
|
self.all_sources = [
|
|
{
|
|
"id": agent_data["source"],
|
|
"retriever": agent_data.get("retriever", "classic"),
|
|
}
|
|
]
|
|
else:
|
|
self.source = {}
|
|
self.all_sources = []
|
|
return
|
|
if "active_docs" in self.data:
|
|
active_docs = self.data["active_docs"]
|
|
if active_docs and active_docs != "default":
|
|
self.source = {"active_docs": active_docs}
|
|
else:
|
|
self.source = {}
|
|
return
|
|
self.source = {}
|
|
self.all_sources = []
|
|
|
|
def _has_active_docs(self) -> bool:
|
|
"""Return True if a real document source is configured for retrieval."""
|
|
active_docs = self.source.get("active_docs") if self.source else None
|
|
if not active_docs:
|
|
return False
|
|
if active_docs == "default":
|
|
return False
|
|
return True
|
|
|
|
def _resolve_agent_id(self) -> Optional[str]:
|
|
"""Resolve agent_id from request, then fall back to conversation context."""
|
|
request_agent_id = self.data.get("agent_id")
|
|
if request_agent_id:
|
|
return str(request_agent_id)
|
|
|
|
if not self.conversation_id or not self.initial_user_id:
|
|
return None
|
|
|
|
try:
|
|
conversation = self.conversation_service.get_conversation(
|
|
self.conversation_id, self.initial_user_id
|
|
)
|
|
except Exception:
|
|
return None
|
|
|
|
if not conversation:
|
|
return None
|
|
|
|
conversation_agent_id = conversation.get("agent_id")
|
|
if conversation_agent_id:
|
|
return str(conversation_agent_id)
|
|
|
|
return None
|
|
|
|
def _configure_agent(self):
|
|
"""Configure the agent based on request data.
|
|
|
|
Unified flow: resolve the effective API key, then extract config once.
|
|
"""
|
|
agent_id = self._resolve_agent_id()
|
|
|
|
self.agent_key, self.is_shared_usage, self.shared_token = self._get_agent_key(
|
|
agent_id, self.initial_user_id
|
|
)
|
|
self.agent_id = str(agent_id) if agent_id else None
|
|
|
|
# Determine the effective API key (explicit > agent-derived)
|
|
effective_key = self.data.get("api_key") or self.agent_key
|
|
|
|
if effective_key:
|
|
self._agent_data = self._get_data_from_api_key(effective_key)
|
|
if self._agent_data.get("_id"):
|
|
self.agent_id = str(self._agent_data.get("_id"))
|
|
|
|
self.agent_config.update(
|
|
{
|
|
"prompt_id": self._agent_data.get("prompt_id", "default"),
|
|
"agent_type": self._agent_data.get("agent_type", settings.AGENT_NAME),
|
|
"user_api_key": effective_key,
|
|
"json_schema": self._agent_data.get("json_schema"),
|
|
"default_model_id": self._agent_data.get("default_model_id", ""),
|
|
"models": self._agent_data.get("models", []),
|
|
"allow_system_prompt_override": self._agent_data.get(
|
|
"allow_system_prompt_override", False
|
|
),
|
|
}
|
|
)
|
|
|
|
# Set identity context
|
|
if self.data.get("api_key"):
|
|
# External API key: use the key owner's identity
|
|
self.initial_user_id = self._agent_data.get("user")
|
|
self.decoded_token = {"sub": self._agent_data.get("user")}
|
|
elif self.is_shared_usage:
|
|
# Shared agent: keep the caller's identity
|
|
pass
|
|
else:
|
|
# Owner using their own agent
|
|
self.decoded_token = {"sub": self._agent_data.get("user")}
|
|
|
|
if self._agent_data.get("workflow"):
|
|
self.agent_config["workflow"] = self._agent_data["workflow"]
|
|
self.agent_config["workflow_owner"] = self._agent_data.get("user")
|
|
else:
|
|
# No API key — default/workflow configuration
|
|
agent_type = settings.AGENT_NAME
|
|
if self.data.get("workflow") and isinstance(
|
|
self.data.get("workflow"), dict
|
|
):
|
|
agent_type = "workflow"
|
|
self.agent_config["workflow"] = self.data["workflow"]
|
|
if isinstance(self.decoded_token, dict):
|
|
self.agent_config["workflow_owner"] = self.decoded_token.get("sub")
|
|
|
|
self.agent_config.update(
|
|
{
|
|
"prompt_id": self.data.get("prompt_id", "default"),
|
|
"agent_type": agent_type,
|
|
"user_api_key": None,
|
|
"json_schema": None,
|
|
"default_model_id": "",
|
|
}
|
|
)
|
|
|
|
def _configure_retriever(self):
|
|
"""Assemble retriever config with precedence: request > agent > default."""
|
|
doc_token_limit = calculate_doc_token_budget(model_id=self.model_id)
|
|
|
|
# Start with defaults
|
|
retriever_name = "classic"
|
|
chunks = 2
|
|
|
|
# Layer agent-level config (if present)
|
|
if self._agent_data:
|
|
if self._agent_data.get("retriever"):
|
|
retriever_name = self._agent_data["retriever"]
|
|
if self._agent_data.get("chunks") is not None:
|
|
try:
|
|
chunks = int(self._agent_data["chunks"])
|
|
except (ValueError, TypeError):
|
|
logger.warning(
|
|
f"Invalid agent chunks value: {self._agent_data['chunks']}, "
|
|
"using default value 2"
|
|
)
|
|
|
|
# Explicit request values win over agent config
|
|
if "retriever" in self.data:
|
|
retriever_name = self.data["retriever"]
|
|
if "chunks" in self.data:
|
|
try:
|
|
chunks = int(self.data["chunks"])
|
|
except (ValueError, TypeError):
|
|
logger.warning(
|
|
f"Invalid request chunks value: {self.data['chunks']}, "
|
|
"using default value 2"
|
|
)
|
|
|
|
self.retriever_config = {
|
|
"retriever_name": retriever_name,
|
|
"chunks": chunks,
|
|
"doc_token_limit": doc_token_limit,
|
|
}
|
|
|
|
# isNoneDoc without an API key forces no retrieval
|
|
api_key = self.data.get("api_key") or self.agent_key
|
|
if not api_key and "isNoneDoc" in self.data and self.data["isNoneDoc"]:
|
|
self.retriever_config["chunks"] = 0
|
|
|
|
def create_retriever(self):
|
|
return RetrieverCreator.create_retriever(
|
|
self.retriever_config["retriever_name"],
|
|
source=self.source,
|
|
chat_history=self.history,
|
|
prompt=get_prompt(self.agent_config["prompt_id"], self.prompts_collection),
|
|
chunks=self.retriever_config["chunks"],
|
|
doc_token_limit=self.retriever_config.get("doc_token_limit", 50000),
|
|
model_id=self.model_id,
|
|
user_api_key=self.agent_config["user_api_key"],
|
|
agent_id=self.agent_id,
|
|
decoded_token=self.decoded_token,
|
|
)
|
|
|
|
def pre_fetch_docs(self, question: str) -> tuple[Optional[str], Optional[list]]:
|
|
"""Pre-fetch documents for template rendering before agent creation"""
|
|
if self.data.get("isNoneDoc", False) and not self.agent_id:
|
|
logger.info("Pre-fetch skipped: isNoneDoc=True")
|
|
return None, None
|
|
if not self._has_active_docs():
|
|
logger.info("Pre-fetch skipped: no active docs configured")
|
|
return None, None
|
|
try:
|
|
retriever = self.create_retriever()
|
|
logger.info(
|
|
f"Pre-fetching docs with chunks={retriever.chunks}, doc_token_limit={retriever.doc_token_limit}"
|
|
)
|
|
docs = retriever.search(question)
|
|
logger.info(f"Pre-fetch retrieved {len(docs) if docs else 0} documents")
|
|
|
|
if not docs:
|
|
logger.info("Pre-fetch: No documents returned from search")
|
|
return None, None
|
|
self.retrieved_docs = docs
|
|
|
|
docs_with_filenames = []
|
|
for doc in docs:
|
|
filename = doc.get("filename") or doc.get("title") or doc.get("source")
|
|
if filename:
|
|
chunk_header = str(filename)
|
|
docs_with_filenames.append(f"{chunk_header}\n{doc['text']}")
|
|
else:
|
|
docs_with_filenames.append(doc["text"])
|
|
docs_together = "\n\n".join(docs_with_filenames)
|
|
|
|
logger.info(f"Pre-fetch docs_together size: {len(docs_together)} chars")
|
|
|
|
return docs_together, docs
|
|
except Exception as e:
|
|
logger.error(f"Failed to pre-fetch docs: {str(e)}", exc_info=True)
|
|
return None, None
|
|
|
|
def pre_fetch_tools(self) -> Optional[Dict[str, Any]]:
|
|
"""Pre-fetch tool data for template rendering before agent creation"""
|
|
if not settings.ENABLE_TOOL_PREFETCH:
|
|
logger.info(
|
|
"Tool pre-fetching disabled globally via ENABLE_TOOL_PREFETCH setting"
|
|
)
|
|
return None
|
|
|
|
if self.data.get("disable_tool_prefetch", False):
|
|
logger.info("Tool pre-fetching disabled for this request")
|
|
return None
|
|
|
|
required_tool_actions = self._get_required_tool_actions()
|
|
filtering_enabled = required_tool_actions is not None
|
|
|
|
try:
|
|
user_tools_collection = self.db["user_tools"]
|
|
user_id = self.initial_user_id or "local"
|
|
|
|
user_tools = list(
|
|
user_tools_collection.find({"user": user_id, "status": True})
|
|
)
|
|
|
|
if not user_tools:
|
|
return None
|
|
|
|
tools_data = {}
|
|
|
|
for tool_doc in user_tools:
|
|
tool_name = tool_doc.get("name")
|
|
tool_id = str(tool_doc.get("_id"))
|
|
|
|
if filtering_enabled:
|
|
required_actions_by_name = required_tool_actions.get(
|
|
tool_name, set()
|
|
)
|
|
required_actions_by_id = required_tool_actions.get(tool_id, set())
|
|
|
|
required_actions = required_actions_by_name | required_actions_by_id
|
|
|
|
if not required_actions:
|
|
continue
|
|
else:
|
|
required_actions = None
|
|
|
|
tool_data = self._fetch_tool_data(tool_doc, required_actions)
|
|
if tool_data:
|
|
tools_data[tool_name] = tool_data
|
|
tools_data[tool_id] = tool_data
|
|
|
|
return tools_data if tools_data else None
|
|
except Exception as e:
|
|
logger.warning(f"Failed to pre-fetch tools: {type(e).__name__}")
|
|
return None
|
|
|
|
def _fetch_tool_data(
|
|
self,
|
|
tool_doc: Dict[str, Any],
|
|
required_actions: Optional[Set[Optional[str]]],
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Fetch and execute tool actions with saved parameters"""
|
|
try:
|
|
from application.agents.tools.tool_manager import ToolManager
|
|
|
|
tool_name = tool_doc.get("name")
|
|
tool_config = tool_doc.get("config", {}).copy()
|
|
tool_config["tool_id"] = str(tool_doc["_id"])
|
|
|
|
tool_manager = ToolManager(config={tool_name: tool_config})
|
|
user_id = self.initial_user_id or "local"
|
|
tool = tool_manager.load_tool(tool_name, tool_config, user_id=user_id)
|
|
|
|
if not tool:
|
|
logger.debug(f"Tool '{tool_name}' failed to load")
|
|
return None
|
|
|
|
tool_actions = tool.get_actions_metadata()
|
|
if not tool_actions:
|
|
logger.debug(f"Tool '{tool_name}' has no actions")
|
|
return None
|
|
|
|
saved_actions = tool_doc.get("actions", [])
|
|
|
|
include_all_actions = required_actions is None or (
|
|
required_actions and None in required_actions
|
|
)
|
|
allowed_actions: Set[str] = (
|
|
{action for action in required_actions if isinstance(action, str)}
|
|
if required_actions
|
|
else set()
|
|
)
|
|
|
|
action_results = {}
|
|
for action_meta in tool_actions:
|
|
action_name = action_meta.get("name")
|
|
if action_name is None:
|
|
continue
|
|
if (
|
|
not include_all_actions
|
|
and allowed_actions
|
|
and action_name not in allowed_actions
|
|
):
|
|
continue
|
|
|
|
try:
|
|
saved_action = None
|
|
for sa in saved_actions:
|
|
if sa.get("name") == action_name:
|
|
saved_action = sa
|
|
break
|
|
|
|
action_params = action_meta.get("parameters", {})
|
|
properties = action_params.get("properties", {})
|
|
|
|
kwargs = {}
|
|
for param_name, param_spec in properties.items():
|
|
if saved_action:
|
|
saved_props = saved_action.get("parameters", {}).get(
|
|
"properties", {}
|
|
)
|
|
if param_name in saved_props:
|
|
param_value = saved_props[param_name].get("value")
|
|
if param_value is not None:
|
|
kwargs[param_name] = param_value
|
|
continue
|
|
|
|
if param_name in tool_config:
|
|
kwargs[param_name] = tool_config[param_name]
|
|
elif "default" in param_spec:
|
|
kwargs[param_name] = param_spec["default"]
|
|
|
|
result = tool.execute_action(action_name, **kwargs)
|
|
action_results[action_name] = result
|
|
except Exception as e:
|
|
logger.debug(
|
|
f"Action '{action_name}' execution failed: {type(e).__name__}"
|
|
)
|
|
continue
|
|
|
|
return action_results if action_results else None
|
|
|
|
except Exception as e:
|
|
logger.debug(f"Tool pre-fetch failed for '{tool_name}': {type(e).__name__}")
|
|
return None
|
|
|
|
def _get_prompt_content(self) -> Optional[str]:
|
|
"""Retrieve and cache the raw prompt content for the current agent configuration."""
|
|
if self._prompt_content is not None:
|
|
return self._prompt_content
|
|
prompt_id = (
|
|
self.agent_config.get("prompt_id")
|
|
if isinstance(self.agent_config, dict)
|
|
else None
|
|
)
|
|
if not prompt_id:
|
|
return None
|
|
try:
|
|
self._prompt_content = get_prompt(prompt_id, self.prompts_collection)
|
|
except ValueError as e:
|
|
logger.debug(f"Invalid prompt ID '{prompt_id}': {str(e)}")
|
|
self._prompt_content = None
|
|
except Exception as e:
|
|
logger.debug(f"Failed to fetch prompt '{prompt_id}': {type(e).__name__}")
|
|
self._prompt_content = None
|
|
return self._prompt_content
|
|
|
|
def _get_required_tool_actions(self) -> Optional[Dict[str, Set[Optional[str]]]]:
|
|
"""Determine which tool actions are referenced in the prompt template"""
|
|
if self._required_tool_actions is not None:
|
|
return self._required_tool_actions
|
|
|
|
prompt_content = self._get_prompt_content()
|
|
if prompt_content is None:
|
|
return None
|
|
|
|
if "{{" not in prompt_content or "}}" not in prompt_content:
|
|
self._required_tool_actions = {}
|
|
return self._required_tool_actions
|
|
|
|
try:
|
|
from application.templates.template_engine import TemplateEngine
|
|
|
|
template_engine = TemplateEngine()
|
|
usages = template_engine.extract_tool_usages(prompt_content)
|
|
self._required_tool_actions = usages
|
|
return self._required_tool_actions
|
|
except Exception as e:
|
|
logger.debug(f"Failed to extract tool usages: {type(e).__name__}")
|
|
self._required_tool_actions = {}
|
|
return self._required_tool_actions
|
|
|
|
def _fetch_memory_tool_data(
|
|
self, tool_doc: Dict[str, Any]
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Fetch memory tool data for pre-injection into prompt"""
|
|
try:
|
|
tool_config = tool_doc.get("config", {}).copy()
|
|
tool_config["tool_id"] = str(tool_doc["_id"])
|
|
|
|
from application.agents.tools.memory import MemoryTool
|
|
|
|
memory_tool = MemoryTool(tool_config, self.initial_user_id)
|
|
|
|
root_view = memory_tool.execute_action("view", path="/")
|
|
|
|
if "Error:" in root_view or not root_view.strip():
|
|
return None
|
|
|
|
return {"root": root_view, "available": True}
|
|
except Exception as e:
|
|
logger.warning(f"Failed to fetch memory tool data: {str(e)}")
|
|
return None
|
|
|
|
def resume_from_tool_actions(
|
|
self,
|
|
tool_actions: list,
|
|
conversation_id: str,
|
|
):
|
|
"""Resume a paused agent from saved continuation state.
|
|
|
|
Loads the pending state from MongoDB, recreates the agent with
|
|
the saved configuration, and returns an agent ready to call
|
|
``gen_continuation()``.
|
|
|
|
Args:
|
|
tool_actions: Client-provided actions (approvals / results).
|
|
conversation_id: The conversation being resumed.
|
|
|
|
Returns:
|
|
Tuple of (agent, messages, tools_dict, pending_tool_calls, tool_actions).
|
|
"""
|
|
from application.api.answer.services.continuation_service import (
|
|
ContinuationService,
|
|
)
|
|
from application.agents.agent_creator import AgentCreator
|
|
from application.agents.tool_executor import ToolExecutor
|
|
from application.llm.handlers.handler_creator import LLMHandlerCreator
|
|
from application.llm.llm_creator import LLMCreator
|
|
|
|
cont_service = ContinuationService()
|
|
state = cont_service.load_state(conversation_id, self.initial_user_id)
|
|
if not state:
|
|
raise ValueError("No pending tool state found for this conversation")
|
|
|
|
messages = state["messages"]
|
|
pending_tool_calls = state["pending_tool_calls"]
|
|
tools_dict = state["tools_dict"]
|
|
tool_schemas = state.get("tool_schemas", [])
|
|
agent_config = state["agent_config"]
|
|
|
|
model_id = agent_config.get("model_id")
|
|
llm_name = agent_config.get("llm_name", settings.LLM_PROVIDER)
|
|
api_key = agent_config.get("api_key")
|
|
user_api_key = agent_config.get("user_api_key")
|
|
agent_id = agent_config.get("agent_id")
|
|
prompt = agent_config.get("prompt", "")
|
|
json_schema = agent_config.get("json_schema")
|
|
retriever_config = agent_config.get("retriever_config")
|
|
|
|
# Recreate dependencies
|
|
system_api_key = api_key or get_api_key_for_provider(llm_name)
|
|
llm = LLMCreator.create_llm(
|
|
llm_name,
|
|
api_key=system_api_key,
|
|
user_api_key=user_api_key,
|
|
decoded_token=self.decoded_token,
|
|
model_id=model_id,
|
|
agent_id=agent_id,
|
|
)
|
|
llm_handler = LLMHandlerCreator.create_handler(llm_name or "default")
|
|
tool_executor = ToolExecutor(
|
|
user_api_key=user_api_key,
|
|
user=self.initial_user_id,
|
|
decoded_token=self.decoded_token,
|
|
)
|
|
tool_executor.conversation_id = conversation_id
|
|
# Restore client tools so they stay available for subsequent LLM calls
|
|
saved_client_tools = state.get("client_tools")
|
|
if saved_client_tools:
|
|
tool_executor.client_tools = saved_client_tools
|
|
# Re-merge into tools_dict (they may have been stripped during serialization)
|
|
tool_executor.merge_client_tools(tools_dict, saved_client_tools)
|
|
|
|
agent_type = agent_config.get("agent_type", "ClassicAgent")
|
|
# Map class names back to agent creator keys
|
|
type_map = {
|
|
"ClassicAgent": "classic",
|
|
"AgenticAgent": "agentic",
|
|
"ResearchAgent": "research",
|
|
"WorkflowAgent": "workflow",
|
|
}
|
|
agent_key = type_map.get(agent_type, "classic")
|
|
|
|
agent_kwargs = {
|
|
"endpoint": "stream",
|
|
"llm_name": llm_name,
|
|
"model_id": model_id,
|
|
"api_key": system_api_key,
|
|
"agent_id": agent_id,
|
|
"user_api_key": user_api_key,
|
|
"prompt": prompt,
|
|
"chat_history": [],
|
|
"decoded_token": self.decoded_token,
|
|
"json_schema": json_schema,
|
|
"llm": llm,
|
|
"llm_handler": llm_handler,
|
|
"tool_executor": tool_executor,
|
|
}
|
|
|
|
if agent_key in ("agentic", "research") and retriever_config:
|
|
agent_kwargs["retriever_config"] = retriever_config
|
|
|
|
agent = AgentCreator.create_agent(agent_key, **agent_kwargs)
|
|
agent.conversation_id = conversation_id
|
|
agent.initial_user_id = self.initial_user_id
|
|
agent.tools = tool_schemas
|
|
|
|
# Store config for the route layer
|
|
self.model_id = model_id
|
|
self.agent_id = agent_id
|
|
self.agent_config["user_api_key"] = user_api_key
|
|
self.conversation_id = conversation_id
|
|
|
|
# Delete state so it can't be replayed
|
|
cont_service.delete_state(conversation_id, self.initial_user_id)
|
|
|
|
return agent, messages, tools_dict, pending_tool_calls, tool_actions
|
|
|
|
def create_agent(
|
|
self,
|
|
docs_together: Optional[str] = None,
|
|
docs: Optional[list] = None,
|
|
tools_data: Optional[Dict[str, Any]] = None,
|
|
):
|
|
"""Create and return the configured agent with rendered prompt"""
|
|
agent_type = self.agent_config["agent_type"]
|
|
|
|
# For agentic agents, swap standard presets for their agentic
|
|
# counterparts (which include search tool instructions instead of
|
|
# {summaries}). Custom / user-provided prompts pass through as-is.
|
|
raw_prompt = self._get_prompt_content()
|
|
if raw_prompt is None:
|
|
prompt_id = self.agent_config.get("prompt_id", "default")
|
|
agentic_presets = {"default", "creative", "strict"}
|
|
if agent_type in ("agentic", "research") and prompt_id in agentic_presets:
|
|
raw_prompt = get_prompt(
|
|
f"agentic_{prompt_id}", self.prompts_collection
|
|
)
|
|
else:
|
|
raw_prompt = get_prompt(prompt_id, self.prompts_collection)
|
|
self._prompt_content = raw_prompt
|
|
|
|
# Allow API callers to override the system prompt when the agent
|
|
# has opted in via allow_system_prompt_override.
|
|
if (
|
|
self.agent_config.get("allow_system_prompt_override", False)
|
|
and self.data.get("system_prompt_override")
|
|
):
|
|
rendered_prompt = self.data["system_prompt_override"]
|
|
else:
|
|
rendered_prompt = self.prompt_renderer.render_prompt(
|
|
prompt_content=raw_prompt,
|
|
user_id=self.initial_user_id,
|
|
request_id=self.data.get("request_id"),
|
|
passthrough_data=self.data.get("passthrough"),
|
|
docs=docs,
|
|
docs_together=docs_together,
|
|
tools_data=tools_data,
|
|
)
|
|
|
|
provider = (
|
|
get_provider_from_model_id(self.model_id)
|
|
if self.model_id
|
|
else settings.LLM_PROVIDER
|
|
)
|
|
system_api_key = get_api_key_for_provider(provider or settings.LLM_PROVIDER)
|
|
|
|
# Create LLM and handler (dependency injection)
|
|
from application.llm.llm_creator import LLMCreator
|
|
from application.llm.handlers.handler_creator import LLMHandlerCreator
|
|
from application.agents.tool_executor import ToolExecutor
|
|
|
|
# Compute backup models: agent's configured models minus the active one
|
|
agent_models = self.agent_config.get("models", [])
|
|
backup_models = [m for m in agent_models if m != self.model_id]
|
|
|
|
llm = LLMCreator.create_llm(
|
|
provider or settings.LLM_PROVIDER,
|
|
api_key=system_api_key,
|
|
user_api_key=self.agent_config["user_api_key"],
|
|
decoded_token=self.decoded_token,
|
|
model_id=self.model_id,
|
|
agent_id=self.agent_id,
|
|
backup_models=backup_models,
|
|
)
|
|
llm_handler = LLMHandlerCreator.create_handler(
|
|
provider if provider else "default"
|
|
)
|
|
|
|
user = self.decoded_token.get("sub") if self.decoded_token else None
|
|
tool_executor = ToolExecutor(
|
|
user_api_key=self.agent_config["user_api_key"],
|
|
user=user,
|
|
decoded_token=self.decoded_token,
|
|
)
|
|
tool_executor.conversation_id = self.conversation_id
|
|
# Pass client-side tools so they get merged in get_tools()
|
|
client_tools = self.data.get("client_tools")
|
|
if client_tools:
|
|
tool_executor.client_tools = client_tools
|
|
|
|
# Base agent kwargs
|
|
agent_kwargs = {
|
|
"endpoint": "stream",
|
|
"llm_name": provider or settings.LLM_PROVIDER,
|
|
"model_id": self.model_id,
|
|
"api_key": system_api_key,
|
|
"agent_id": self.agent_id,
|
|
"user_api_key": self.agent_config["user_api_key"],
|
|
"prompt": rendered_prompt,
|
|
"chat_history": self.history,
|
|
"retrieved_docs": self.retrieved_docs,
|
|
"decoded_token": self.decoded_token,
|
|
"attachments": self.attachments,
|
|
"json_schema": self.agent_config.get("json_schema"),
|
|
"compressed_summary": self.compressed_summary,
|
|
"llm": llm,
|
|
"llm_handler": llm_handler,
|
|
"tool_executor": tool_executor,
|
|
}
|
|
|
|
# Type-specific kwargs
|
|
if agent_type in ("agentic", "research"):
|
|
agent_kwargs["retriever_config"] = {
|
|
"source": self.source,
|
|
"retriever_name": self.retriever_config.get(
|
|
"retriever_name", "classic"
|
|
),
|
|
"chunks": self.retriever_config.get("chunks", 2),
|
|
"doc_token_limit": self.retriever_config.get(
|
|
"doc_token_limit", 50000
|
|
),
|
|
"model_id": self.model_id,
|
|
"user_api_key": self.agent_config["user_api_key"],
|
|
"agent_id": self.agent_id,
|
|
"llm_name": provider or settings.LLM_PROVIDER,
|
|
"api_key": system_api_key,
|
|
"decoded_token": self.decoded_token,
|
|
}
|
|
|
|
elif agent_type == "workflow":
|
|
workflow_config = self.agent_config.get("workflow")
|
|
if isinstance(workflow_config, str):
|
|
agent_kwargs["workflow_id"] = workflow_config
|
|
elif isinstance(workflow_config, dict):
|
|
agent_kwargs["workflow"] = workflow_config
|
|
workflow_owner = self.agent_config.get("workflow_owner")
|
|
if workflow_owner:
|
|
agent_kwargs["workflow_owner"] = workflow_owner
|
|
|
|
agent = AgentCreator.create_agent(agent_type, **agent_kwargs)
|
|
|
|
agent.conversation_id = self.conversation_id
|
|
agent.initial_user_id = self.initial_user_id
|
|
|
|
return agent
|