Compare commits
10 Commits
582bf27deb
...
feat/proje
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
cc0e258e8c | ||
|
|
12e203e63d | ||
|
|
ffcd7390f0 | ||
|
|
91e880f9d4 | ||
|
|
7d47ca54be | ||
|
|
956fa88853 | ||
|
|
fb2f59ccea | ||
|
|
56dbb7f4cd | ||
|
|
506f517851 | ||
|
|
520c186991 |
168
app/agents/folder_agent.py
Normal file
168
app/agents/folder_agent.py
Normal file
@@ -0,0 +1,168 @@
|
||||
"""Scoped file-read and search tools for the project folder feature."""
|
||||
from __future__ import annotations
|
||||
|
||||
from langchain_core.tools import tool
|
||||
|
||||
from app.core.folder_indexer import _extract_docx_text, _extract_pdf_text
|
||||
from app.core.ws_context import execute_on_client
|
||||
|
||||
# Cap returned slice size to keep tool output under control.
|
||||
_MAX_RETURN_CHARS = 50_000
|
||||
_MAX_SEARCH_MATCHES = 20
|
||||
|
||||
|
||||
def _is_unsafe_path(rel: str) -> bool:
|
||||
if not rel:
|
||||
return True
|
||||
norm = rel.replace("\\", "/")
|
||||
if norm.startswith("/"):
|
||||
return True
|
||||
# Windows drive letter
|
||||
if len(rel) >= 2 and rel[1] == ":":
|
||||
return True
|
||||
parts = norm.split("/")
|
||||
return ".." in parts
|
||||
|
||||
|
||||
async def _fetch_file(project_id: str, relative_path: str, offset: int, length: int) -> dict:
|
||||
"""Return the raw Electron tool_result dict for a file read."""
|
||||
return await execute_on_client(
|
||||
action="read_project_folder_file",
|
||||
data={
|
||||
"projectId": project_id,
|
||||
"relativePath": relative_path,
|
||||
"offset": offset,
|
||||
"length": length,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _decode(result: dict) -> tuple[str, str, int]:
|
||||
"""Decode a tool_result into (text, kind, total_size). For pdf/docx,
|
||||
extracts text from base64. For images, returns a placeholder string.
|
||||
For text, content is already a sliced utf-8 string.
|
||||
"""
|
||||
kind = result.get("kind", "text")
|
||||
content = result.get("content", "") or ""
|
||||
total = int(result.get("totalSize", 0) or 0)
|
||||
if kind == "image":
|
||||
return ("[Image file — cannot be navigated as text. See manifest summary.]", kind, total)
|
||||
if kind == "pdf":
|
||||
return (_extract_pdf_text(content), kind, total)
|
||||
if kind == "docx":
|
||||
return (_extract_docx_text(content), kind, total)
|
||||
return (content, kind, total)
|
||||
|
||||
|
||||
@tool
|
||||
async def read_project_folder_file(
|
||||
project_id: str,
|
||||
relative_path: str,
|
||||
offset: int = 0,
|
||||
length: int = _MAX_RETURN_CHARS,
|
||||
) -> str:
|
||||
"""Read a slice of a file inside the project's linked folder.
|
||||
|
||||
Args:
|
||||
project_id: project ID.
|
||||
relative_path: path relative to the linked folder root.
|
||||
offset: char offset to start reading from (0 = beginning).
|
||||
length: max chars to return. Default 50000. Use smaller values to save tokens.
|
||||
|
||||
Returns text content slice with a header showing position. Header tells you
|
||||
when more content is available; call again with the suggested next offset.
|
||||
|
||||
For PDF / DOCX files the backend extracts text first, then applies offset/length
|
||||
on the extracted text. For images returns a placeholder; navigate with the
|
||||
manifest summary instead.
|
||||
"""
|
||||
if _is_unsafe_path(relative_path):
|
||||
return "Access denied"
|
||||
|
||||
result = await _fetch_file(project_id, relative_path, offset, length)
|
||||
text, kind, total_size = _decode(result)
|
||||
|
||||
if not text and kind in ("missing", "error"):
|
||||
return f"File not found or unreadable: {relative_path}"
|
||||
|
||||
if kind in ("pdf", "docx"):
|
||||
# Backend extracted full text — apply offset/length on chars.
|
||||
sliced = text[offset:offset + length]
|
||||
slice_end = min(offset + length, len(text))
|
||||
header = (
|
||||
f"[file={relative_path} kind={kind} offset={offset} end={slice_end} "
|
||||
f"totalChars={len(text)}]"
|
||||
)
|
||||
if slice_end < len(text):
|
||||
header += f"\n[More content available — call again with offset={slice_end}.]"
|
||||
return header + "\n" + sliced
|
||||
|
||||
if kind == "text":
|
||||
slice_end = offset + len(text)
|
||||
header = (
|
||||
f"[file={relative_path} kind=text offset={offset} end={slice_end} "
|
||||
f"totalBytes={total_size}]"
|
||||
)
|
||||
if slice_end < total_size:
|
||||
header += f"\n[More content available — call again with offset={slice_end}.]"
|
||||
return header + "\n" + text
|
||||
|
||||
# image or unknown
|
||||
return text
|
||||
|
||||
|
||||
@tool
|
||||
async def search_project_folder_file(
|
||||
project_id: str,
|
||||
relative_path: str,
|
||||
query: str,
|
||||
context_lines: int = 3,
|
||||
) -> str:
|
||||
"""Search a project folder file for a query string (case-insensitive substring).
|
||||
|
||||
Args:
|
||||
project_id: project ID.
|
||||
relative_path: path relative to the linked folder root.
|
||||
query: text to search for.
|
||||
context_lines: number of lines of context around each match (default 3).
|
||||
|
||||
Returns matching line ranges with surrounding context and 1-based line numbers.
|
||||
Capped at 20 matches; if more exist the header shows the total.
|
||||
|
||||
Works on text, code, markdown, PDF (extracted), and DOCX (extracted).
|
||||
Images and binary files are not searchable.
|
||||
"""
|
||||
if _is_unsafe_path(relative_path):
|
||||
return "Access denied"
|
||||
if not query:
|
||||
return "Empty query."
|
||||
|
||||
# For text we still need full file; pass length=very large.
|
||||
result = await _fetch_file(project_id, relative_path, offset=0, length=10_000_000)
|
||||
text, kind, _ = _decode(result)
|
||||
|
||||
if not text and kind in ("missing", "error"):
|
||||
return f"File not found or unreadable: {relative_path}"
|
||||
if kind == "image":
|
||||
return "Cannot search inside images."
|
||||
|
||||
lines = text.splitlines()
|
||||
q = query.lower()
|
||||
matches = [i for i, line in enumerate(lines) if q in line.lower()]
|
||||
if not matches:
|
||||
return f"No matches for '{query}' in {relative_path}."
|
||||
|
||||
shown = matches[:_MAX_SEARCH_MATCHES]
|
||||
snippets: list[str] = []
|
||||
for i in shown:
|
||||
start = max(0, i - context_lines)
|
||||
end = min(len(lines), i + context_lines + 1)
|
||||
block = "\n".join(f"{n + 1:5d}: {lines[n]}" for n in range(start, end))
|
||||
snippets.append(block)
|
||||
|
||||
header = f"[file={relative_path} matches={len(matches)} showing={len(shown)} query='{query}']"
|
||||
body = "\n---\n".join(snippets)
|
||||
return header + "\n" + body
|
||||
|
||||
|
||||
FOLDER_TOOLS = [read_project_folder_file, search_project_folder_file]
|
||||
@@ -228,11 +228,13 @@ async def _handle_home_request(
|
||||
request_id = frame.get("request_id") or str(uuid4())
|
||||
message: str = frame.get("message", "")
|
||||
session_id: str = frame.get("session_id") or str(uuid4())
|
||||
project_id: str | None = frame.get("project_id") or frame.get("projectId") or None
|
||||
logger.info(
|
||||
"device_ws: home_request_start user=%s req=%s session=%s msg=%s",
|
||||
"device_ws: home_request_start user=%s req=%s session=%s project=%s msg=%s",
|
||||
user_id,
|
||||
request_id,
|
||||
session_id,
|
||||
project_id,
|
||||
message[:200],
|
||||
)
|
||||
|
||||
@@ -257,7 +259,7 @@ async def _handle_home_request(
|
||||
set_client_executor(executor)
|
||||
response_chunks: list[str] = []
|
||||
try:
|
||||
event_stream = run_home_stream(user_id, message, context)
|
||||
event_stream = run_home_stream(user_id, message, context, project_id=project_id)
|
||||
formatter = StreamFormatter(request_id=request_id)
|
||||
async for ws_frame in formatter.format(event_stream):
|
||||
await websocket.send_text(ws_frame.model_dump_json())
|
||||
@@ -454,10 +456,11 @@ async def _handle_task_brief_request(
|
||||
request_id = frame.get("request_id") or str(uuid4())
|
||||
session_id = frame.get("session_id") or str(uuid4())
|
||||
task_id: str = frame.get("task_id") or frame.get("taskId") or ""
|
||||
project_id: str | None = frame.get("project_id") or frame.get("projectId") or None
|
||||
|
||||
logger.info(
|
||||
"device_ws: task_brief_request_start user=%s req=%s task=%s [cache_miss]",
|
||||
user_id, request_id, task_id,
|
||||
"device_ws: task_brief_request_start user=%s req=%s task=%s project=%s [cache_miss]",
|
||||
user_id, request_id, task_id, project_id,
|
||||
)
|
||||
|
||||
if not task_id:
|
||||
@@ -486,7 +489,7 @@ async def _handle_task_brief_request(
|
||||
response_chunks: list[str] = []
|
||||
|
||||
try:
|
||||
event_stream = run_task_brief_research_stream(user_id, task_id, context)
|
||||
event_stream = run_task_brief_research_stream(user_id, task_id, context, project_id=project_id)
|
||||
formatter = StreamFormatter(request_id=request_id)
|
||||
async for ws_frame in formatter.format(event_stream):
|
||||
if ws_frame.type == "stream_text": # type: ignore[union-attr]
|
||||
@@ -595,9 +598,9 @@ async def _handle_index_session_start(
|
||||
frame: dict,
|
||||
) -> None:
|
||||
"""Register a new folder index session. No response sent — client is declaring intent."""
|
||||
session_id: str = frame.get("sessionId") or frame.get("session_id", "")
|
||||
session_id: str = frame.get("sessionId") or frame.get("session_id") or ""
|
||||
project_id: str | None = frame.get("projectId") or frame.get("project_id")
|
||||
total: int = int(frame.get("totalFiles", 0))
|
||||
total: int = int(frame.get("totalFiles") or frame.get("total_files") or 0)
|
||||
|
||||
if not session_id:
|
||||
logger.warning("device_ws: index_session_start missing sessionId user=%s", user_id)
|
||||
@@ -621,7 +624,7 @@ async def _handle_index_session_cancel(
|
||||
frame: dict,
|
||||
) -> None:
|
||||
"""Mark a session as cancelled and emit index_session_done(cancelled)."""
|
||||
session_id: str = frame.get("sessionId") or frame.get("session_id", "")
|
||||
session_id: str = frame.get("sessionId") or frame.get("session_id") or ""
|
||||
session = _index_sessions.get(session_id)
|
||||
if session:
|
||||
session["cancelled"] = True
|
||||
@@ -651,7 +654,7 @@ async def _handle_index_file_batch(
|
||||
from app.billing.tier_manager import tier_manager # noqa: PLC0415
|
||||
from app.billing.quota import add_token_usage # noqa: PLC0415
|
||||
|
||||
session_id: str = frame.get("sessionId") or frame.get("session_id", "")
|
||||
session_id: str = frame.get("sessionId") or frame.get("session_id") or ""
|
||||
files: list[dict] = frame.get("files", [])
|
||||
|
||||
session = _index_sessions.get(session_id)
|
||||
@@ -667,11 +670,12 @@ async def _handle_index_file_batch(
|
||||
if session.get("cancelled"):
|
||||
return
|
||||
|
||||
rel_path: str = file_info.get("relPath", "")
|
||||
kind: str = file_info.get("kind", "text")
|
||||
content: str = file_info.get("content", "")
|
||||
ext: str = file_info.get("ext", "")
|
||||
mime: str = file_info.get("mime", "application/octet-stream")
|
||||
# Electron's toSnakeCase converts payload keys, so accept both forms.
|
||||
rel_path: str = file_info.get("relPath") or file_info.get("rel_path") or ""
|
||||
kind: str = file_info.get("kind") or "text"
|
||||
content: str = file_info.get("content") or ""
|
||||
ext: str = file_info.get("ext") or ""
|
||||
mime: str = file_info.get("mime") or "application/octet-stream"
|
||||
name: str = rel_path.split("/")[-1] or rel_path
|
||||
|
||||
try:
|
||||
|
||||
@@ -21,6 +21,7 @@ from app.core.deep_agent import (
|
||||
_relational_memory_injection,
|
||||
_run_single_agent_stream,
|
||||
_trace_id_from_context,
|
||||
build_brief_multi_project_manifest,
|
||||
)
|
||||
from app.core.langfuse_client import compile_prompt, get_prompt_or_fallback
|
||||
|
||||
@@ -159,6 +160,8 @@ async def run_home_brief(
|
||||
Yields (event_type, data) tuples identical to _run_single_agent_stream.
|
||||
Do NOT post-process output through _normalize_tagged_list_lines.
|
||||
"""
|
||||
from app.agents.folder_agent import FOLDER_TOOLS
|
||||
|
||||
trace_id = _trace_id_from_context(context)
|
||||
today = date.today().isoformat()
|
||||
language = _resolve_language(context)
|
||||
@@ -171,7 +174,10 @@ async def run_home_brief(
|
||||
if today not in system_prompt:
|
||||
system_prompt += f"\nToday is {today}."
|
||||
|
||||
tools = _build_read_tools(user_id, trace_id)
|
||||
brief_manifest = await build_brief_multi_project_manifest()
|
||||
system_prompt = system_prompt + ("\n\n" + brief_manifest if brief_manifest else "")
|
||||
|
||||
tools = [*_build_read_tools(user_id, trace_id), *FOLDER_TOOLS]
|
||||
async for event in _run_single_agent_stream(
|
||||
user_id=user_id,
|
||||
system_prompt=system_prompt,
|
||||
|
||||
@@ -60,6 +60,93 @@ def _language_instruction(context: dict[str, Any]) -> str:
|
||||
f"All your output text must be written in {lang}."
|
||||
)
|
||||
|
||||
MANIFEST_TOKEN_BUDGET = 3000 # rough budget for <linked_folder> block
|
||||
|
||||
|
||||
def format_folder_manifest(manifest: dict | None) -> str:
|
||||
"""Format a folder manifest into the <linked_folder> block.
|
||||
|
||||
Truncates by mtime DESC if estimated tokens exceed MANIFEST_TOKEN_BUDGET.
|
||||
Returns empty string if manifest is None or has no files.
|
||||
"""
|
||||
if not manifest or not manifest.get("files"):
|
||||
return ""
|
||||
files = list(manifest["files"])
|
||||
files.sort(key=lambda f: f.get("mtimeMs", 0), reverse=True)
|
||||
|
||||
header = (
|
||||
f"<linked_folder>\npath: {manifest.get('folderPath', '?')} "
|
||||
f"({len(files)} files, scanned {manifest.get('lastScannedAt', '?')})\nfiles:\n"
|
||||
)
|
||||
footer_template = "… {} more files omitted, use read_project_folder_file to access by path\n</linked_folder>"
|
||||
|
||||
char_budget = MANIFEST_TOKEN_BUDGET * 4 # ~4 chars/token
|
||||
body = ""
|
||||
included = 0
|
||||
for f in files:
|
||||
line = f"- /{f['relPath']} [{f.get('kind','text')}] {f.get('summary','')}\n"
|
||||
if len(header) + len(body) + len(line) + len(footer_template.format(0)) > char_budget:
|
||||
break
|
||||
body += line
|
||||
included += 1
|
||||
omitted = len(files) - included
|
||||
if omitted > 0:
|
||||
return header + body + footer_template.format(omitted)
|
||||
return header + body + "</linked_folder>"
|
||||
|
||||
|
||||
async def _fetch_project_manifest(project_id: str) -> dict | None:
|
||||
"""Fetch manifest from Electron via execute_on_client. Returns None if unlinked or error."""
|
||||
from app.core.ws_context import execute_on_client
|
||||
try:
|
||||
result = await execute_on_client(
|
||||
action="read_project_folder_manifest",
|
||||
data={"projectId": project_id},
|
||||
)
|
||||
if not result or not result.get("folderPath"):
|
||||
return None
|
||||
return result
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def build_brief_multi_project_manifest() -> str:
|
||||
"""Build a compact multi-project manifest for the daily brief agent.
|
||||
|
||||
Calls execute_on_client('list_projects_with_folder_manifests') and keeps
|
||||
the top 5 most-recently-modified files per project.
|
||||
"""
|
||||
try:
|
||||
result = await execute_on_client(
|
||||
action="list_projects_with_folder_manifests",
|
||||
data={},
|
||||
)
|
||||
except Exception:
|
||||
return ""
|
||||
projects = (result or {}).get("projects") or []
|
||||
if not projects:
|
||||
return ""
|
||||
blocks: list[str] = ["<linked_folders>"]
|
||||
any_entry = False
|
||||
for p in projects:
|
||||
all_files = p.get("files", []) or []
|
||||
files = sorted(all_files, key=lambda f: f.get("mtimeMs", 0), reverse=True)[:5]
|
||||
blocks.append(f"project: {p.get('projectName','?')} [{p.get('projectId','?')}]")
|
||||
blocks.append(f" path: {p.get('folderPath','?')} (scanned {p.get('lastScannedAt','?')})")
|
||||
if not all_files:
|
||||
blocks.append(" (no indexed files yet — folder is linked but empty or unscanned)")
|
||||
else:
|
||||
for f in files:
|
||||
blocks.append(f" - /{f['relPath']} [{f.get('kind','text')}] {f.get('summary','')}")
|
||||
if len(all_files) > 5:
|
||||
blocks.append(f" … {len(all_files) - 5} more files (use read_project_folder_file by relPath)")
|
||||
any_entry = True
|
||||
if not any_entry:
|
||||
return ""
|
||||
blocks.append("</linked_folders>")
|
||||
return "\n".join(blocks)
|
||||
|
||||
|
||||
def _datetime_context_injection(context: dict[str, Any]) -> str:
|
||||
"""Build a comprehensive DATE CONTEXT block with pre-computed ms-epoch boundaries for common ranges."""
|
||||
fp = context.get("format_prefs")
|
||||
@@ -1328,9 +1415,26 @@ async def run_home_stream(
|
||||
user_id: str,
|
||||
message: str,
|
||||
context: dict[str, Any],
|
||||
project_id: str | None = None,
|
||||
) -> AsyncGenerator[tuple[str, Any], None]:
|
||||
from app.agents.folder_agent import FOLDER_TOOLS
|
||||
|
||||
prepared_context = await _prepare_context(message, context)
|
||||
system_prompt, langfuse_prompt = _build_system_prompt("home_system", _HOME_SYSTEM_PROMPT, prepared_context)
|
||||
|
||||
manifest_block = ""
|
||||
if project_id:
|
||||
manifest = await _fetch_project_manifest(project_id)
|
||||
manifest_block = format_folder_manifest(manifest)
|
||||
if not manifest_block:
|
||||
# No specific project context — surface all linked folders so the agent
|
||||
# can answer questions like "tell me about project X" using its files.
|
||||
manifest_block = await build_brief_multi_project_manifest()
|
||||
system_prompt = system_prompt + ("\n\n" + manifest_block if manifest_block else "")
|
||||
|
||||
trace_id = _trace_id_from_context(prepared_context)
|
||||
tools = [*_all_tools_for_user(user_id, trace_id), *FOLDER_TOOLS]
|
||||
|
||||
text_chunks: list[str] = []
|
||||
async for event in _run_single_agent_stream(
|
||||
user_id=user_id,
|
||||
@@ -1339,6 +1443,7 @@ async def run_home_stream(
|
||||
context=prepared_context,
|
||||
langfuse_prompt=langfuse_prompt,
|
||||
agent_name="home-agent",
|
||||
tools=tools,
|
||||
conversation_history=context.get("conversation_history"),
|
||||
):
|
||||
event_type, data = event
|
||||
@@ -1421,6 +1526,7 @@ async def run_task_brief_research_stream(
|
||||
user_id: str,
|
||||
task_id: str,
|
||||
context: dict[str, Any],
|
||||
project_id: str | None = None,
|
||||
) -> AsyncGenerator[tuple[str, Any], None]:
|
||||
"""Stage-1 executive assistant: deep research for one task.
|
||||
|
||||
@@ -1428,8 +1534,10 @@ async def run_task_brief_research_stream(
|
||||
The final concatenated text may contain a ``<canvas kind="...">...</canvas>`` block
|
||||
which the WS handler strips and emits as a ``canvas_draft`` mutation.
|
||||
"""
|
||||
from app.agents.folder_agent import FOLDER_TOOLS
|
||||
|
||||
prepared_context = await _prepare_context(f"task:{task_id}", context)
|
||||
tools = _brief_research_tools(user_id, _trace_id_from_context(prepared_context))
|
||||
tools = [*_brief_research_tools(user_id, _trace_id_from_context(prepared_context)), *FOLDER_TOOLS]
|
||||
|
||||
# Inject task_id so the agent knows what to look up first.
|
||||
research_message = (
|
||||
@@ -1446,6 +1554,12 @@ async def run_task_brief_research_stream(
|
||||
prepared_context,
|
||||
)
|
||||
|
||||
manifest_block = ""
|
||||
if project_id:
|
||||
manifest = await _fetch_project_manifest(project_id)
|
||||
manifest_block = format_folder_manifest(manifest)
|
||||
system_prompt = system_prompt + ("\n\n" + manifest_block if manifest_block else "")
|
||||
|
||||
async for event in _run_single_agent_stream(
|
||||
user_id=user_id,
|
||||
system_prompt=system_prompt,
|
||||
|
||||
@@ -12,6 +12,7 @@ from docx import Document as DocxDocument
|
||||
from app.core.langfuse_client import (
|
||||
compile_prompt,
|
||||
extract_usage,
|
||||
get_langfuse,
|
||||
get_prompt_or_fallback,
|
||||
)
|
||||
from app.core.llm import get_llm
|
||||
@@ -55,7 +56,7 @@ async def _llm_vision(messages: list) -> object:
|
||||
return await llm.ainvoke(messages)
|
||||
|
||||
|
||||
async def summarize_image(*, image_b64: str, mime: str) -> IndexResult:
|
||||
async def summarize_image(*, image_b64: str, mime: str, file_name: str | None = None) -> IndexResult:
|
||||
"""Return a compact summary of an image file using vision.
|
||||
|
||||
Parameters
|
||||
@@ -64,6 +65,8 @@ async def summarize_image(*, image_b64: str, mime: str) -> IndexResult:
|
||||
Base64-encoded image bytes.
|
||||
mime:
|
||||
MIME type of the image, e.g. ``"image/png"``.
|
||||
file_name:
|
||||
Optional file name, attached to the Langfuse trace as input metadata.
|
||||
"""
|
||||
template, prompt_obj = get_prompt_or_fallback("folder_file_summary_image", _IMAGE_FALLBACK)
|
||||
messages = [
|
||||
@@ -73,6 +76,19 @@ async def summarize_image(*, image_b64: str, mime: str) -> IndexResult:
|
||||
{"type": "image_url", "image_url": {"url": f"data:{mime};base64,{image_b64}"}},
|
||||
]),
|
||||
]
|
||||
lf = get_langfuse()
|
||||
if lf is not None:
|
||||
with lf.start_as_current_observation(
|
||||
as_type="generation",
|
||||
name="folder-summarize-image",
|
||||
model="gpt-4o-mini",
|
||||
prompt=prompt_obj,
|
||||
input={"file_name": file_name, "mime": mime},
|
||||
) as gen:
|
||||
response = await _llm_vision(messages)
|
||||
usage = extract_usage(response)
|
||||
gen.update(output=response.content, usage_details=usage)
|
||||
else:
|
||||
response = await _llm_vision(messages)
|
||||
usage = extract_usage(response)
|
||||
summary = (response.content or "").strip()[:500]
|
||||
@@ -98,6 +114,19 @@ async def summarize_text(*, content: str, ext: str, name: str) -> IndexResult:
|
||||
SystemMessage(content=compiled),
|
||||
HumanMessage(content="Summarise this file."),
|
||||
]
|
||||
lf = get_langfuse()
|
||||
if lf is not None:
|
||||
with lf.start_as_current_observation(
|
||||
as_type="generation",
|
||||
name="folder-summarize-text",
|
||||
model="gpt-4o-mini",
|
||||
prompt=prompt_obj,
|
||||
input={"file_name": name, "ext": ext, "content_chars": len(truncated)},
|
||||
) as gen:
|
||||
response = await _llm_text(messages)
|
||||
usage = extract_usage(response)
|
||||
gen.update(output=response.content, usage_details=usage)
|
||||
else:
|
||||
response = await _llm_text(messages)
|
||||
usage = extract_usage(response)
|
||||
summary = (response.content or "").strip()[:500]
|
||||
|
||||
139
tests/test_folder_agent_tool.py
Normal file
139
tests/test_folder_agent_tool.py
Normal file
@@ -0,0 +1,139 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from app.agents.folder_agent import (
|
||||
read_project_folder_file,
|
||||
search_project_folder_file,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
async def test_happy_path():
|
||||
with patch(
|
||||
"app.agents.folder_agent.execute_on_client",
|
||||
new=AsyncMock(return_value={"content": "file body", "kind": "text", "totalSize": 9}),
|
||||
):
|
||||
out = await read_project_folder_file.ainvoke({"project_id": "p1", "relative_path": "docs/x.md"})
|
||||
assert "file body" in out
|
||||
assert "kind=text" in out
|
||||
|
||||
|
||||
async def test_traversal_rejected():
|
||||
out = await read_project_folder_file.ainvoke({"project_id": "p1", "relative_path": "../../etc/passwd"})
|
||||
assert out == "Access denied"
|
||||
|
||||
|
||||
async def test_absolute_path_rejected():
|
||||
out = await read_project_folder_file.ainvoke({"project_id": "p1", "relative_path": "C:\\Windows\\foo"})
|
||||
assert out == "Access denied"
|
||||
|
||||
|
||||
async def test_missing_file():
|
||||
with patch(
|
||||
"app.agents.folder_agent.execute_on_client",
|
||||
new=AsyncMock(return_value={"content": "", "kind": "missing", "totalSize": 0}),
|
||||
):
|
||||
out = await read_project_folder_file.ainvoke({"project_id": "p1", "relative_path": "ghost.md"})
|
||||
assert "not found" in out.lower()
|
||||
|
||||
|
||||
async def test_pagination_signals_more_available():
|
||||
# Electron returned the first slice, totalSize larger than slice length.
|
||||
with patch(
|
||||
"app.agents.folder_agent.execute_on_client",
|
||||
new=AsyncMock(return_value={"content": "first chunk", "kind": "text", "totalSize": 1000}),
|
||||
):
|
||||
out = await read_project_folder_file.ainvoke({
|
||||
"project_id": "p1",
|
||||
"relative_path": "big.txt",
|
||||
"offset": 0,
|
||||
"length": 11,
|
||||
})
|
||||
assert "first chunk" in out
|
||||
assert "More content available" in out
|
||||
assert "offset=11" in out
|
||||
|
||||
|
||||
async def test_pdf_extracted_then_sliced(monkeypatch):
|
||||
from app.agents import folder_agent
|
||||
monkeypatch.setattr(folder_agent, "_extract_pdf_text", lambda b: "ABC " * 100)
|
||||
with patch(
|
||||
"app.agents.folder_agent.execute_on_client",
|
||||
new=AsyncMock(return_value={"content": "JVBERi0xLg==", "kind": "pdf", "totalSize": 12}),
|
||||
):
|
||||
out = await read_project_folder_file.ainvoke({
|
||||
"project_id": "p1",
|
||||
"relative_path": "doc.pdf",
|
||||
"offset": 0,
|
||||
"length": 8,
|
||||
})
|
||||
assert "kind=pdf" in out
|
||||
assert "ABC ABC " in out
|
||||
assert "More content available" in out
|
||||
|
||||
|
||||
async def test_image_returns_placeholder():
|
||||
with patch(
|
||||
"app.agents.folder_agent.execute_on_client",
|
||||
new=AsyncMock(return_value={"content": "iVBORw0K", "kind": "image", "totalSize": 1024}),
|
||||
):
|
||||
out = await read_project_folder_file.ainvoke({"project_id": "p1", "relative_path": "logo.png"})
|
||||
assert "image" in out.lower()
|
||||
|
||||
|
||||
async def test_search_finds_match_with_context():
|
||||
body = "alpha\nbeta\nthe needle is here\ngamma\ndelta"
|
||||
with patch(
|
||||
"app.agents.folder_agent.execute_on_client",
|
||||
new=AsyncMock(return_value={"content": body, "kind": "text", "totalSize": len(body)}),
|
||||
):
|
||||
out = await search_project_folder_file.ainvoke({
|
||||
"project_id": "p1",
|
||||
"relative_path": "log.txt",
|
||||
"query": "needle",
|
||||
"context_lines": 1,
|
||||
})
|
||||
assert "needle" in out
|
||||
assert "matches=1" in out
|
||||
# Context lines included
|
||||
assert "beta" in out
|
||||
assert "gamma" in out
|
||||
|
||||
|
||||
async def test_search_no_match():
|
||||
with patch(
|
||||
"app.agents.folder_agent.execute_on_client",
|
||||
new=AsyncMock(return_value={"content": "nothing here", "kind": "text", "totalSize": 12}),
|
||||
):
|
||||
out = await search_project_folder_file.ainvoke({
|
||||
"project_id": "p1",
|
||||
"relative_path": "x.txt",
|
||||
"query": "zzz",
|
||||
})
|
||||
assert "No matches" in out
|
||||
|
||||
|
||||
async def test_search_rejects_traversal():
|
||||
out = await search_project_folder_file.ainvoke({
|
||||
"project_id": "p1",
|
||||
"relative_path": "../etc/passwd",
|
||||
"query": "root",
|
||||
})
|
||||
assert out == "Access denied"
|
||||
|
||||
|
||||
async def test_search_image_rejected():
|
||||
with patch(
|
||||
"app.agents.folder_agent.execute_on_client",
|
||||
new=AsyncMock(return_value={"content": "b64data", "kind": "image", "totalSize": 100}),
|
||||
):
|
||||
out = await search_project_folder_file.ainvoke({
|
||||
"project_id": "p1",
|
||||
"relative_path": "logo.png",
|
||||
"query": "anything",
|
||||
})
|
||||
assert "Cannot search" in out
|
||||
69
tests/test_manifest_injection.py
Normal file
69
tests/test_manifest_injection.py
Normal file
@@ -0,0 +1,69 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from app.core.deep_agent import format_folder_manifest, MANIFEST_TOKEN_BUDGET
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
def test_format_folder_manifest_basic():
|
||||
manifest = {
|
||||
"folderPath": "D:\\Acme",
|
||||
"lastScannedAt": "2h ago",
|
||||
"files": [
|
||||
{"relPath": "briefs/kickoff.md", "kind": "text", "summary": "Kickoff notes; scope and deadlines."},
|
||||
{"relPath": "logos/logo-v3.png", "kind": "image", "summary": "Final logo on white."},
|
||||
],
|
||||
}
|
||||
out = format_folder_manifest(manifest)
|
||||
assert "<linked_folder>" in out
|
||||
assert "/briefs/kickoff.md" in out or "briefs/kickoff.md" in out
|
||||
assert "[text]" in out
|
||||
assert "[image]" in out
|
||||
|
||||
|
||||
def test_format_folder_manifest_truncates_past_budget():
|
||||
files = [
|
||||
{"relPath": f"f{i}.md", "kind": "text", "summary": "x" * 100, "mtimeMs": i}
|
||||
for i in range(2000)
|
||||
]
|
||||
out = format_folder_manifest({"folderPath": "p", "lastScannedAt": "now", "files": files})
|
||||
assert "more files omitted" in out
|
||||
# Rough token check
|
||||
assert len(out) // 4 < MANIFEST_TOKEN_BUDGET + 200
|
||||
|
||||
|
||||
def test_format_folder_manifest_null_returns_empty():
|
||||
assert format_folder_manifest(None) == ""
|
||||
assert format_folder_manifest({"files": []}) == ""
|
||||
|
||||
|
||||
async def test_brief_multi_project_manifest_top_5_per_project():
|
||||
fake_response = [
|
||||
{
|
||||
"projectId": "p1", "projectName": "Acme", "folderPath": "/a",
|
||||
"lastScannedAt": "now",
|
||||
"files": [
|
||||
{"relPath": f"f{i}.md", "kind": "text", "summary": "s", "mtimeMs": i}
|
||||
for i in range(10)
|
||||
],
|
||||
},
|
||||
{
|
||||
"projectId": "p2", "projectName": "Beta", "folderPath": "/b",
|
||||
"lastScannedAt": "now",
|
||||
"files": [{"relPath": "x.md", "kind": "text", "summary": "s", "mtimeMs": 1}],
|
||||
},
|
||||
]
|
||||
with patch(
|
||||
"app.core.deep_agent.execute_on_client",
|
||||
new=AsyncMock(return_value={"projects": fake_response}),
|
||||
):
|
||||
from app.core.deep_agent import build_brief_multi_project_manifest
|
||||
out = await build_brief_multi_project_manifest()
|
||||
# Project 1 has 10 files, only top 5 by mtimeMs should appear
|
||||
assert out.count("[p1]") <= 5
|
||||
# Project 2 has 1 file, must appear
|
||||
assert "[p2]" in out or "Beta" in out
|
||||
Reference in New Issue
Block a user