fe0dd038ee
- Langfuse SDK v4: fix prompt-to-trace linking (as_type=generation) - tracing: compile_prompt with Langfuse managed prompt fallback - journey: remove journey CLI subcommand (keep only interactive) - LLM: add service-specific llm modules for batch-agent and chat - gitignore: exclude eval private test data - config: add LANGFUSE settings to shared config
396 lines
14 KiB
Python
396 lines
14 KiB
Python
"""Chatbot Journey — guided conversation to build an agent prompt_template.
|
||
|
||
Adapted for Batch Agent Service: imports from app.agents.filesystem_agent
|
||
and app.llm instead of monolith paths. Session state is in-memory (could
|
||
be moved to Redis for horizontal scaling in the future).
|
||
|
||
Journey flow:
|
||
1. Redis consumer dispatches ``journey_start`` with basic agent config.
|
||
2. Server creates an in-memory session, runs the setup LLM with
|
||
file-system tools to explore the directory, returns first question.
|
||
3. ``journey_message`` frames drive the conversation.
|
||
4. After 3-5 turns the LLM emits PROMPT_TEMPLATE_START / _END block.
|
||
5. Server parses the block and returns ``journey_reply`` with ``done=True``.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import time
|
||
import uuid
|
||
from dataclasses import dataclass, field
|
||
from typing import Any
|
||
|
||
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
|
||
|
||
from app.agents.filesystem_agent import FILESYSTEM_TOOLS
|
||
from shared.llm import get_llm
|
||
import app.tracing as tracing
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# ── Session TTL ───────────────────────────────────────────────────────────
|
||
|
||
_SESSION_TTL_SECONDS: int = 1800 # 30 minutes
|
||
|
||
# Sentinel strings used to delimit the LLM-produced prompt_template.
|
||
_TEMPLATE_START = "PROMPT_TEMPLATE_START"
|
||
_TEMPLATE_END = "PROMPT_TEMPLATE_END"
|
||
|
||
_MIN_TURNS_BEFORE_NUDGE: int = 3
|
||
_MAX_TURNS: int = 15
|
||
_MAX_TOOL_STEPS: int = 6
|
||
|
||
# ── In-memory session store ───────────────────────────────────────────────
|
||
|
||
|
||
@dataclass
|
||
class JourneySession:
|
||
session_id: str
|
||
user_id: str
|
||
agent_type: str # "local" | "cloud"
|
||
directory: str
|
||
data_types: list[str]
|
||
history: list[dict[str, Any]] = field(default_factory=list)
|
||
system_prompt: str = ""
|
||
created_at: float = field(default_factory=time.monotonic)
|
||
|
||
def is_expired(self) -> bool:
|
||
return (time.monotonic() - self.created_at) > _SESSION_TTL_SECONDS
|
||
|
||
|
||
# session_id → session
|
||
_sessions: dict[str, JourneySession] = {}
|
||
|
||
|
||
def get_journey_session(session_id: str, user_id: str) -> JourneySession | None:
|
||
"""Retrieve session; return None on missing, expired, or wrong owner."""
|
||
s = _sessions.get(session_id)
|
||
if s is None or s.is_expired():
|
||
_sessions.pop(session_id, None)
|
||
return None
|
||
if s.user_id != user_id:
|
||
return None
|
||
return s
|
||
|
||
|
||
# ── System prompt builder ─────────────────────────────────────────────────
|
||
|
||
_SYSTEM_PROMPT_TEMPLATE = """\
|
||
You are a friendly assistant helping a freelancer configure a data-extraction agent.
|
||
Your job is to understand exactly what data the user wants to extract from their
|
||
local directory and produce a concise prompt_template that a separate AI will use
|
||
as its instruction set.
|
||
|
||
You have access to file-system tools to explore the user's directory:
|
||
- list_directory: to see folder structure
|
||
- read_file_content: to peek at file contents
|
||
- get_file_metadata: to check file info
|
||
|
||
The user's configured directory is: {directory}
|
||
Target data types: {data_types}
|
||
|
||
IMPORTANT — project assignment is handled automatically. You MUST NOT ask the user
|
||
about projects, projectId, or how to link records to projects. Never include
|
||
projectId logic or project creation instructions in the generated prompt_template.
|
||
|
||
Start by exploring the directory to understand its structure. Then ask concise,
|
||
focused questions one at a time. Cover only the topics relevant to the target
|
||
data types listed above:
|
||
|
||
1. Content type and format — confirmed by your exploration.
|
||
2. For TASKS (if in scope): field mapping for title, status, priority, content,
|
||
dueDate (where is the date found? what's the fallback when absent?),
|
||
and assignee (is there a person name to assign?).
|
||
3. For NOTES when TASKS are also in scope: note vs task distinction —
|
||
what makes something a note rather than a task?
|
||
4. For TIMELINES (if in scope): the date source — what marks a milestone or event?
|
||
5. Exclusions and special handling applicable to the target data types.
|
||
|
||
Keep asking focused questions until you are at least 90% confident. Then stop and
|
||
output the final prompt_template immediately, wrapped between these exact markers
|
||
on their own lines:
|
||
|
||
{template_start}
|
||
<the complete extraction prompt here>
|
||
{template_end}
|
||
|
||
The prompt_template must be concise (bullet points, ~15–25 lines maximum).
|
||
Specify only:
|
||
- Scope: what files/content qualify and what entity types to create.
|
||
- Field mapping rules per entity type (camelCase fields: title, status, priority,
|
||
dueDate, content, assignee, etc.).
|
||
- dueDate rule (if tasks in scope): source and fallback behaviour.
|
||
- Note vs task rule (if both in scope): the criterion that separates them.
|
||
- Timeline date rule (if timelines in scope): what constitutes a timeline event.
|
||
- Exclusion/filtering rules.
|
||
- 2–3 concrete mapping examples based on what you discovered.
|
||
|
||
{existing_section}Begin by exploring the directory, then ask your first question.\
|
||
"""
|
||
|
||
|
||
def _build_system_prompt(
|
||
directory: str,
|
||
data_types: list[str],
|
||
existing_template: str | None = None,
|
||
) -> str:
|
||
existing_section = (
|
||
f"\nThe user already has the following prompt_template — refine it based on their answers:\n"
|
||
f"---\n{existing_template}\n---\n"
|
||
if existing_template
|
||
else ""
|
||
)
|
||
# Use Langfuse compile_prompt ({{variable}} syntax) with Python .format() fallback
|
||
return tracing.compile_prompt(
|
||
"journey_system",
|
||
fallback=_SYSTEM_PROMPT_TEMPLATE,
|
||
variables={
|
||
"directory": directory,
|
||
"data_types": ", ".join(data_types),
|
||
"existing_section": existing_section,
|
||
},
|
||
)
|
||
|
||
|
||
# ── Template extraction ───────────────────────────────────────────────────
|
||
|
||
|
||
def _extract_template(text: str) -> str | None:
|
||
"""Return the text between PROMPT_TEMPLATE_START and PROMPT_TEMPLATE_END, or None."""
|
||
if _TEMPLATE_START not in text or _TEMPLATE_END not in text:
|
||
return None
|
||
start_idx = text.index(_TEMPLATE_START) + len(_TEMPLATE_START)
|
||
end_idx = text.index(_TEMPLATE_END)
|
||
return text[start_idx:end_idx].strip() or None
|
||
|
||
|
||
# ── LLM call with tool support ───────────────────────────────────────────
|
||
|
||
|
||
def _as_text(content: Any) -> str:
|
||
if content is None:
|
||
return ""
|
||
if isinstance(content, str):
|
||
return content
|
||
if isinstance(content, list):
|
||
parts: list[str] = []
|
||
for item in content:
|
||
if isinstance(item, str):
|
||
parts.append(item)
|
||
elif isinstance(item, dict):
|
||
text = item.get("text")
|
||
if isinstance(text, str):
|
||
parts.append(text)
|
||
return "".join(parts)
|
||
return str(content)
|
||
|
||
|
||
async def _call_llm_with_tools(
|
||
system_prompt: str,
|
||
history: list[dict[str, Any]],
|
||
tools: list[Any],
|
||
langfuse_handler: Any | None = None,
|
||
) -> str:
|
||
"""Build LangChain messages from history and invoke the LLM with tools.
|
||
|
||
Handles tool-calling loops: if the LLM calls tools, execute them and
|
||
continue until a final text response is produced.
|
||
"""
|
||
messages: list[Any] = [SystemMessage(content=system_prompt)]
|
||
for turn in history:
|
||
if turn["role"] == "user":
|
||
messages.append(HumanMessage(content=turn["content"]))
|
||
else:
|
||
messages.append(AIMessage(content=turn["content"]))
|
||
|
||
callbacks = [langfuse_handler] if langfuse_handler else None
|
||
llm = get_llm(model=None, temperature=0.4, callbacks=callbacks)
|
||
llm_with_tools = llm.bind_tools(tools)
|
||
tool_map = {tool_def.name: tool_def for tool_def in tools}
|
||
|
||
for _ in range(_MAX_TOOL_STEPS):
|
||
response: AIMessage = await llm_with_tools.ainvoke(messages)
|
||
messages.append(response)
|
||
|
||
if not response.tool_calls:
|
||
return _as_text(response.content)
|
||
|
||
for call in response.tool_calls:
|
||
call_name = str(call.get("name", ""))
|
||
call_args = call.get("args", {})
|
||
logger.info(
|
||
"journey: tool_call name=%s args=%s",
|
||
call_name,
|
||
json.dumps(call_args, ensure_ascii=True)[:500],
|
||
)
|
||
|
||
tool_fn = tool_map.get(call_name)
|
||
if tool_fn is None:
|
||
tool_output = f"Unknown tool: {call_name}"
|
||
else:
|
||
tool_output = await tool_fn.ainvoke(call_args)
|
||
|
||
logger.info(
|
||
"journey: tool_result name=%s output=%s",
|
||
call_name,
|
||
str(tool_output)[:800],
|
||
)
|
||
messages.append(ToolMessage(content=str(tool_output), tool_call_id=call["id"]))
|
||
|
||
# Fallback: exceeded max tool steps.
|
||
final = await llm.ainvoke(messages)
|
||
return _as_text(final.content)
|
||
|
||
|
||
# ── Journey handlers (called from redis_consumer) ────────────────────────
|
||
|
||
|
||
async def handle_journey_start(
|
||
user_id: str,
|
||
frame: dict[str, Any],
|
||
*,
|
||
langfuse_handler: Any | None = None,
|
||
) -> dict[str, Any]:
|
||
"""Handle a ``journey_start`` request.
|
||
|
||
Creates a session, runs the setup LLM with directory exploration,
|
||
and returns the ``journey_reply`` payload.
|
||
"""
|
||
agent_type = frame.get("agent_type", "local")
|
||
directory = frame.get("directory", "")
|
||
data_types = frame.get("data_types", [])
|
||
existing_template = frame.get("existing_template")
|
||
|
||
session_id = frame.get("session_id") or str(uuid.uuid4())
|
||
system_prompt = _build_system_prompt(directory, data_types, existing_template)
|
||
|
||
session = JourneySession(
|
||
session_id=session_id,
|
||
user_id=user_id,
|
||
agent_type=agent_type,
|
||
directory=directory,
|
||
data_types=data_types,
|
||
system_prompt=system_prompt,
|
||
)
|
||
|
||
seed_history: list[dict[str, Any]] = [
|
||
{"role": "user", "content": "Hi, I'm ready to set up my agent. Please explore my directory and ask me your first question."},
|
||
]
|
||
ai_reply = await _call_llm_with_tools(
|
||
system_prompt=system_prompt,
|
||
history=seed_history,
|
||
tools=list(FILESYSTEM_TOOLS),
|
||
langfuse_handler=langfuse_handler,
|
||
)
|
||
|
||
session.history.extend(seed_history)
|
||
session.history.append({"role": "assistant", "content": ai_reply})
|
||
_sessions[session_id] = session
|
||
|
||
logger.info(
|
||
"journey: session %s started for user %s (directory=%s)",
|
||
session_id,
|
||
user_id,
|
||
directory,
|
||
)
|
||
|
||
prompt_template = _extract_template(ai_reply)
|
||
done = prompt_template is not None
|
||
|
||
display_message = ai_reply
|
||
if done:
|
||
display_message = (
|
||
ai_reply[: ai_reply.index(_TEMPLATE_START)].strip()
|
||
or "Here is your agent configuration. You can save it or continue refining."
|
||
)
|
||
_sessions.pop(session_id, None)
|
||
|
||
return {
|
||
"type": "journey_reply",
|
||
"session_id": session_id,
|
||
"message": display_message,
|
||
"done": done,
|
||
"prompt_template": prompt_template,
|
||
}
|
||
|
||
|
||
async def handle_journey_message(
|
||
user_id: str,
|
||
frame: dict[str, Any],
|
||
*,
|
||
langfuse_handler: Any | None = None,
|
||
) -> dict[str, Any]:
|
||
"""Handle a ``journey_message`` request.
|
||
|
||
Appends the user message, calls the LLM, and returns the
|
||
``journey_reply`` payload.
|
||
"""
|
||
session_id = frame.get("session_id", "")
|
||
message = frame.get("message", "")
|
||
|
||
session = get_journey_session(session_id, user_id)
|
||
if session is None:
|
||
return {
|
||
"type": "journey_reply",
|
||
"session_id": session_id,
|
||
"message": "Journey session not found or expired. Please start a new setup.",
|
||
"done": True,
|
||
"prompt_template": None,
|
||
}
|
||
|
||
session.history.append({"role": "user", "content": message})
|
||
|
||
ai_reply = await _call_llm_with_tools(
|
||
system_prompt=session.system_prompt,
|
||
history=session.history,
|
||
tools=list(FILESYSTEM_TOOLS),
|
||
langfuse_handler=langfuse_handler,
|
||
)
|
||
|
||
session.history.append({"role": "assistant", "content": ai_reply})
|
||
|
||
prompt_template = _extract_template(ai_reply)
|
||
done = prompt_template is not None
|
||
|
||
if not done:
|
||
turns = sum(1 for t in session.history if t["role"] == "user")
|
||
if turns >= _MAX_TURNS:
|
||
nudge_content = (
|
||
"[System: You have enough information. Please generate the final "
|
||
f"prompt_template now, wrapped in {_TEMPLATE_START} / {_TEMPLATE_END} markers.]"
|
||
)
|
||
session.history.append({"role": "user", "content": nudge_content})
|
||
|
||
nudge_reply = await _call_llm_with_tools(
|
||
system_prompt=session.system_prompt,
|
||
history=session.history,
|
||
tools=list(FILESYSTEM_TOOLS),
|
||
langfuse_handler=langfuse_handler,
|
||
)
|
||
session.history.append({"role": "assistant", "content": nudge_reply})
|
||
|
||
prompt_template = _extract_template(nudge_reply)
|
||
if prompt_template is not None:
|
||
done = True
|
||
ai_reply = nudge_reply
|
||
|
||
display_message = ai_reply
|
||
if done:
|
||
display_message = (
|
||
ai_reply[: ai_reply.index(_TEMPLATE_START)].strip()
|
||
if _TEMPLATE_START in ai_reply
|
||
else "Here is your agent configuration. You can save it or continue refining."
|
||
)
|
||
_sessions.pop(session_id, None)
|
||
logger.info("journey: session %s completed for user %s", session_id, user_id)
|
||
|
||
return {
|
||
"type": "journey_reply",
|
||
"session_id": session_id,
|
||
"message": display_message,
|
||
"done": done,
|
||
"prompt_template": prompt_template,
|
||
}
|