From 0f639faa66210a0551c41724be24d0416f8e8fa1 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Thu, 9 Jul 2026 12:46:10 +0000 Subject: [PATCH] feat: Implement Cognito Agent Phase 1 (Loop + Tools) - Added SSE endpoint `POST /api/agent/loop` for agent execution. - Implemented core agent loop with tool support. - Added four initial tools: `read`, `write`, `edit`, and `bash`. - Implemented security measures: `ProjectTrustStore`, `PROTECTED_FILES`, and `sudo` rejection. - Enhanced `BackendClient` and `BackendRouter` with tool-calling support and failover. - Refactored uncertainty calculation to `app/core/uncertainty.py`. - Added unit tests for tools, trust store, and agent loop. - Updated documentation and added `AGENTS.md`. Co-authored-by: Axlfc <14998495+Axlfc@users.noreply.github.com> --- .../cognito-backend/AGENTS.md | 5 + .../cognito-backend/README.md | 21 +++ .../app/api/routes/ai_agents.py | 69 +++++++++ .../app/api/routes/openai_compat.py | 40 +---- .../cognito-backend/app/core/__init__.py | 0 .../cognito-backend/app/core/agent_loop.py | 142 +++++++++++++++++ .../cognito-backend/app/core/events.py | 31 ++++ .../cognito-backend/app/core/project_trust.py | 36 +++++ .../app/core/protected_files.py | 7 + .../app/core/resource_loader.py | 44 ++++++ .../app/core/tools/__init__.py | 0 .../cognito-backend/app/core/tools/base.py | 21 +++ .../app/core/tools/bash_tool.py | 55 +++++++ .../app/core/tools/edit_tool.py | 61 ++++++++ .../app/core/tools/read_tool.py | 45 ++++++ .../app/core/tools/write_tool.py | 46 ++++++ .../cognito-backend/app/core/uncertainty.py | 39 +++++ .../app/services/backend_client.py | 146 +++++++++++++++++- .../app/services/backend_router.py | 48 +++++- .../cognito-backend/tests/test_agent_loop.py | 90 +++++++++++ .../tests/test_project_trust.py | 17 ++ .../cognito-backend/tests/test_tools.py | 95 ++++++++++++ 22 files changed, 1016 insertions(+), 42 deletions(-) create mode 100644 very-simplified-stack/cognito-backend/AGENTS.md create mode 100644 very-simplified-stack/cognito-backend/app/core/__init__.py create mode 100644 very-simplified-stack/cognito-backend/app/core/agent_loop.py create mode 100644 very-simplified-stack/cognito-backend/app/core/events.py create mode 100644 very-simplified-stack/cognito-backend/app/core/project_trust.py create mode 100644 very-simplified-stack/cognito-backend/app/core/protected_files.py create mode 100644 very-simplified-stack/cognito-backend/app/core/resource_loader.py create mode 100644 very-simplified-stack/cognito-backend/app/core/tools/__init__.py create mode 100644 very-simplified-stack/cognito-backend/app/core/tools/base.py create mode 100644 very-simplified-stack/cognito-backend/app/core/tools/bash_tool.py create mode 100644 very-simplified-stack/cognito-backend/app/core/tools/edit_tool.py create mode 100644 very-simplified-stack/cognito-backend/app/core/tools/read_tool.py create mode 100644 very-simplified-stack/cognito-backend/app/core/tools/write_tool.py create mode 100644 very-simplified-stack/cognito-backend/app/core/uncertainty.py create mode 100644 very-simplified-stack/cognito-backend/tests/test_agent_loop.py create mode 100644 very-simplified-stack/cognito-backend/tests/test_project_trust.py create mode 100644 very-simplified-stack/cognito-backend/tests/test_tools.py diff --git a/very-simplified-stack/cognito-backend/AGENTS.md b/very-simplified-stack/cognito-backend/AGENTS.md new file mode 100644 index 0000000..04c6887 --- /dev/null +++ b/very-simplified-stack/cognito-backend/AGENTS.md @@ -0,0 +1,5 @@ +# AGENTS.md — cognito-backend + +- No modifiques `app/services/backend_registry.py` ni `app/services/semantic_orchestrator.py` sin aprobación explícita del usuario. +- No modifiques `/v1/chat/completions` (comportamiento estable, usado en producción). +- Los tests de regresión de `/api/agent` y `/v1/chat/completions` deben seguir pasando. diff --git a/very-simplified-stack/cognito-backend/README.md b/very-simplified-stack/cognito-backend/README.md index 268decf..276ff3c 100644 --- a/very-simplified-stack/cognito-backend/README.md +++ b/very-simplified-stack/cognito-backend/README.md @@ -59,10 +59,31 @@ Settings are loaded in the following order of priority: - `app/api/routes/openai_compat.py`: Core streaming and uncertainty calculation logic. - `app/services/backend_client.py`: Unified async client for Ollama and OpenAI backends. +- `app/core/agent_loop.py`: Turning text generation into a tool-executing agent loop. - `test-voice-api.ps1`: The main PowerShell profile script containing `cog` and `cogt`. - `Install-CognitoProfile.ps1`: Installer for the PowerShell environment. - `config.example.json`: Template for the user configuration file. +## 🤖 Cognito Agent (Fase 1) + +El backend ahora incluye soporte para agentes capaces de ejecutar herramientas (tools). + +### Endpoints +- `POST /api/agent/loop`: Endpoint SSE que ejecuta un bucle de razonamiento y ejecución de herramientas. + - **Body**: `{ "messages": [...], "cwd": "path/to/repo", "model_params": {} }` + - **Eventos**: `text_delta`, `tool_call`, `tool_result`, `done`, `error`. + +### Herramientas Disponibles +1. `read`: Lee archivos del sistema (restringido al `cwd`). +2. `write`: Crea o sobrescribe archivos (requiere `trust`). +3. `edit`: Edición basada en búsqueda y reemplazo único (requiere `trust`). +4. `bash`: Ejecución de comandos en el workspace (requiere `trust`, sin `sudo`). + +### Seguridad y Trust +- **Protected Files**: Ciertos archivos críticos (`auth.js`, etc.) nunca pueden ser modificados. +- **Project Trust**: Las herramientas de escritura y ejecución requieren que el directorio haya sido marcado como confiable. +- **AGENTS.md**: Si existe en el raíz del `cwd`, se inyecta automáticamente como contexto del sistema. + ## 🧪 Testing To test the uncertainty features: diff --git a/very-simplified-stack/cognito-backend/app/api/routes/ai_agents.py b/very-simplified-stack/cognito-backend/app/api/routes/ai_agents.py index 363fdb0..6f27413 100644 --- a/very-simplified-stack/cognito-backend/app/api/routes/ai_agents.py +++ b/very-simplified-stack/cognito-backend/app/api/routes/ai_agents.py @@ -1,6 +1,19 @@ +import json +from typing import List, Optional, Any, Dict from fastapi import APIRouter, HTTPException +from fastapi.responses import StreamingResponse +from pydantic import BaseModel from app.models.ai import AIRequest, AIResponse from app.services.reasoning_engine import reasoning_engine +from app.services.backend_router import backend_router +from app.core.agent_loop import agent_loop +from app.core.tools.base import ToolContext +from app.core.tools.read_tool import ReadTool +from app.core.tools.write_tool import WriteTool +from app.core.tools.edit_tool import EditTool +from app.core.tools.bash_tool import BashTool +from app.core.project_trust import ProjectTrustStore +from app.core.resource_loader import ResourceLoader import logging # Configure logging @@ -9,6 +22,11 @@ router = APIRouter() +class AgentLoopRequest(BaseModel): + messages: List[Dict[str, Any]] + cwd: str + model_params: Optional[Dict[str, Any]] = None + @router.post("/agent", response_model=AIResponse) async def run_ai_agent(request: AIRequest): """ @@ -27,3 +45,54 @@ async def run_ai_agent(request: AIRequest): status_code=500, detail="An unexpected error occurred while processing the AI request." ) + +@router.post("/agent/loop") +async def run_agent_loop(request: AgentLoopRequest): + """ + Agent Loop endpoint (SSE). + """ + logger.info(f"Starting agent loop in {request.cwd}") + + loader = ResourceLoader(request.cwd) + trust_store = ProjectTrustStore() + + context = ToolContext( + cwd=request.cwd, + trusted=trust_store.is_trusted(request.cwd), + protected_files=loader.get_effective_protected_files() + ) + + tools = [ReadTool(), WriteTool(), EditTool(), BashTool()] + + # Inyectar AGENTS.md en el system prompt si existe + agents_md = loader.discover_agents_md() + messages = list(request.messages) + if agents_md: + system_msg = next((m for m in messages if m["role"] == "system"), None) + if system_msg: + system_msg["content"] += f"\n\nContext from AGENTS.md:\n{agents_md}" + else: + messages.insert(0, {"role": "system", "content": f"Context from AGENTS.md:\n{agents_md}"}) + + async def event_generator(): + try: + async for event in agent_loop( + messages=messages, + tools=tools, + context=context, + backend_router=backend_router, + model_params=request.model_params + ): + yield f"data: {event.model_dump_json()}\n\n" + except Exception as e: + logger.error(f"Error in agent loop generator: {e}", exc_info=True) + yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n" + + return StreamingResponse( + event_generator(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", + } + ) diff --git a/very-simplified-stack/cognito-backend/app/api/routes/openai_compat.py b/very-simplified-stack/cognito-backend/app/api/routes/openai_compat.py index ff21810..396d43e 100644 --- a/very-simplified-stack/cognito-backend/app/api/routes/openai_compat.py +++ b/very-simplified-stack/cognito-backend/app/api/routes/openai_compat.py @@ -44,6 +44,7 @@ from app.services.backend_registry import BACKENDS_BY_PRIORITY, BackendType from app.services.backend_client import BackendClient from app.models.ai import AIRequest +from app.core.uncertainty import compute_uncertainty logger = logging.getLogger(__name__) @@ -259,45 +260,6 @@ async def chat_completions(req: ChatCompletionRequest): ) -# ══════════════════════════════════════════════════════════════════════════════ -# Uncertainty Calculation -# ══════════════════════════════════════════════════════════════════════════════ - -def compute_uncertainty(logprob_data: Any) -> Optional[float]: - """ - Shannon entropy of the top-k distribution, normalized to [0, 1]. - Expects Ollama-style logprobs data. - """ - if not logprob_data: - return None - - top_logprobs = {} - if isinstance(logprob_data, list) and logprob_data: - # Ollama returns a list of token info, usually just one for the current token - entry = logprob_data[0] - for candidate in entry.get("top_logprobs", []): - top_logprobs[candidate["token"]] = candidate["logprob"] - elif isinstance(logprob_data, dict): - # Already parsed or different format - top_logprobs = logprob_data - - if not top_logprobs: - return None - - try: - probs = [math.exp(lp) for lp in top_logprobs.values()] - total = sum(probs) - if total == 0: - return 0.0 - probs = [p / total for p in probs] - entropy = -sum(p * math.log2(p) for p in probs if p > 0) - max_entropy = math.log2(len(probs)) if len(probs) > 1 else 1.0 - return entropy / max_entropy if max_entropy > 0 else 0.0 - except Exception as e: - logger.error("[Uncertainty] Error computing entropy: %s", e) - return None - - # ══════════════════════════════════════════════════════════════════════════════ # Streaming logic # ══════════════════════════════════════════════════════════════════════════════ diff --git a/very-simplified-stack/cognito-backend/app/core/__init__.py b/very-simplified-stack/cognito-backend/app/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/very-simplified-stack/cognito-backend/app/core/agent_loop.py b/very-simplified-stack/cognito-backend/app/core/agent_loop.py new file mode 100644 index 0000000..6865bfd --- /dev/null +++ b/very-simplified-stack/cognito-backend/app/core/agent_loop.py @@ -0,0 +1,142 @@ +import json +import logging +import uuid +from typing import AsyncIterator, Dict, List, Optional, Any + +from app.core.events import ( + AgentEvent, TextDeltaEvent, ToolCallEvent, ToolResultEvent, DoneEvent, ErrorEvent +) +from app.core.tools.base import AgentTool, ToolContext, ToolResult +from app.core.uncertainty import compute_uncertainty + +logger = logging.getLogger(__name__) + +async def agent_loop( + messages: List[Dict[str, Any]], + tools: List[AgentTool], + context: ToolContext, + backend_router, + model_params: Optional[Dict[str, Any]] = None, + max_turns: int = 10, +) -> AsyncIterator[AgentEvent]: + """ + Main Agent Loop: + 1. Call backend with tools. + 2. Stream text deltas. + 3. Handle tool calls: execute tool, emit ToolResultEvent, add to messages. + 4. Repeat until end_turn or max_turns. + """ + + current_messages = list(messages) + # Convert AgentTool list to JSON Schema format for the backend + tools_schema = [ + { + "type": "function", + "function": { + "name": t.name, + "description": t.description, + "parameters": t.parameters_schema, + } + } + for t in tools + ] + + turn = 0 + while turn < max_turns: + turn += 1 + logger.info(f"Starting agent turn {turn}/{max_turns}") + + assistant_content = "" + tool_calls_to_exec = [] + + try: + async for chunk in backend_router.generate_with_tools(current_messages, tools_schema, model_params): + # Text delta + if chunk.get("token"): + token = chunk["token"] + assistant_content += token + uncertainty = compute_uncertainty(chunk.get("logprobs")) + yield TextDeltaEvent(content=token, uncertainty=uncertainty) + + # Tool calls (Ollama style: array in one chunk; OpenAI style: might be fragmented) + # For Phase 1 we assume they arrive complete enough to parse + if chunk.get("tool_calls"): + for tc in chunk["tool_calls"]: + # Normalize format between Ollama and OpenAI + # Ollama: {'function': {'name': ..., 'arguments': {...}}} + # OpenAI: {'id': ..., 'type': 'function', 'function': {...}} + + tc_id = tc.get("id") or f"call_{uuid.uuid4().hex[:8]}" + fn = tc.get("function", {}) + name = fn.get("name") + args = fn.get("arguments") + + if isinstance(args, str): + try: + args = json.loads(args) + except: + pass + + tool_calls_to_exec.append({ + "id": tc_id, + "name": name, + "arguments": args + }) + + yield ToolCallEvent( + tool_call_id=tc_id, + tool_name=name, + arguments=args if isinstance(args, dict) else {"raw": args} + ) + + # Add assistant message to history + # If there were tool calls, we need to add the assistant message with tool_calls + assistant_message = {"role": "assistant", "content": assistant_content} + if tool_calls_to_exec: + # OpenAI format expects tool_calls here + assistant_message["tool_calls"] = [ + { + "id": tc["id"], + "type": "function", + "function": {"name": tc["name"], "arguments": json.dumps(tc["arguments"])} + } + for tc in tool_calls_to_exec + ] + current_messages.append(assistant_message) + + if not tool_calls_to_exec: + # No more tools requested, we are done + yield DoneEvent(stop_reason="end_turn") + return + + # Execute tools + for tc in tool_calls_to_exec: + tool = next((t for t in tools if t.name == tc["name"]), None) + if not tool: + result = ToolResult(is_error=True, output=f"Tool '{tc['name']}' not found.") + else: + logger.info(f"Executing tool {tool.name} with args {tc['arguments']}") + result = await tool.execute(tc["arguments"], context) + + yield ToolResultEvent( + tool_call_id=tc["id"], + tool_name=tc["name"], + output=result.output, + is_error=result.is_error + ) + + # Add tool result to history + current_messages.append({ + "role": "tool", + "tool_call_id": tc["id"], + "name": tc["name"], + "content": result.output + }) + + except Exception as e: + logger.error(f"Error in agent loop turn {turn}: {e}", exc_info=True) + yield ErrorEvent(message=str(e)) + yield DoneEvent(stop_reason="error", error_message=str(e)) + return + + yield DoneEvent(stop_reason="max_turns") diff --git a/very-simplified-stack/cognito-backend/app/core/events.py b/very-simplified-stack/cognito-backend/app/core/events.py new file mode 100644 index 0000000..5510700 --- /dev/null +++ b/very-simplified-stack/cognito-backend/app/core/events.py @@ -0,0 +1,31 @@ +from typing import Literal, Optional, Any, Union +from pydantic import BaseModel + +class TextDeltaEvent(BaseModel): + type: Literal["text_delta"] = "text_delta" + content: str + uncertainty: Optional[float] = None + +class ToolCallEvent(BaseModel): + type: Literal["tool_call"] = "tool_call" + tool_call_id: str + tool_name: str + arguments: dict[str, Any] + +class ToolResultEvent(BaseModel): + type: Literal["tool_result"] = "tool_result" + tool_call_id: str + tool_name: str + output: str + is_error: bool = False + +class DoneEvent(BaseModel): + type: Literal["done"] = "done" + stop_reason: Literal["end_turn", "tool_use", "max_turns", "error", "aborted"] + error_message: Optional[str] = None + +class ErrorEvent(BaseModel): + type: Literal["error"] = "error" + message: str + +AgentEvent = Union[TextDeltaEvent, ToolCallEvent, ToolResultEvent, DoneEvent, ErrorEvent] diff --git a/very-simplified-stack/cognito-backend/app/core/project_trust.py b/very-simplified-stack/cognito-backend/app/core/project_trust.py new file mode 100644 index 0000000..cad8afd --- /dev/null +++ b/very-simplified-stack/cognito-backend/app/core/project_trust.py @@ -0,0 +1,36 @@ +import json +import os +from pathlib import Path +from typing import Dict + +class ProjectTrustStore: + def __init__(self, store_path: Path | None = None): + self.store_path = store_path or (Path.home() / ".cognito" / "trust.json") + self._ensure_dir() + + def _ensure_dir(self): + self.store_path.parent.mkdir(parents=True, exist_ok=True) + + def _load(self) -> Dict[str, bool]: + if not self.store_path.exists(): + return {} + try: + with open(self.store_path, "r") as f: + return json.load(f) + except (json.JSONDecodeError, IOError): + return {} + + def _save(self, data: Dict[str, bool]): + with open(self.store_path, "w") as f: + json.dump(data, f, indent=2) + + def is_trusted(self, repo_path: str) -> bool: + path = os.path.realpath(repo_path) + data = self._load() + return data.get(path, False) + + def set_trusted(self, repo_path: str, trusted: bool) -> None: + path = os.path.realpath(repo_path) + data = self._load() + data[path] = trusted + self._save(data) diff --git a/very-simplified-stack/cognito-backend/app/core/protected_files.py b/very-simplified-stack/cognito-backend/app/core/protected_files.py new file mode 100644 index 0000000..62fbafa --- /dev/null +++ b/very-simplified-stack/cognito-backend/app/core/protected_files.py @@ -0,0 +1,7 @@ +# Mirror of the "auth is sacred" convention used in hypenosys.github.io +PROTECTED_FILES: set[str] = { + "assets/javascript/auth.js", + "assets/javascript/github-api.js", + "assets/javascript/dashboard-data.js", + "dashboard.html", +} diff --git a/very-simplified-stack/cognito-backend/app/core/resource_loader.py b/very-simplified-stack/cognito-backend/app/core/resource_loader.py new file mode 100644 index 0000000..21cfef6 --- /dev/null +++ b/very-simplified-stack/cognito-backend/app/core/resource_loader.py @@ -0,0 +1,44 @@ +import os +import re +from typing import Set +from app.core.protected_files import PROTECTED_FILES + +class ResourceLoader: + def __init__(self, cwd: str): + self.cwd = os.path.realpath(cwd) + + def discover_agents_md(self) -> str: + """ + Looks for AGENTS.md in the root of cwd. + If it exists, returns its content. + """ + agents_md_path = os.path.join(self.cwd, "AGENTS.md") + if os.path.exists(agents_md_path): + try: + with open(agents_md_path, "r") as f: + return f.read() + except Exception: + return "" + return "" + + def get_effective_protected_files(self) -> Set[str]: + """ + Combines default PROTECTED_FILES with any listed in AGENTS.md. + Expects a format like '- Protected: path/to/file' or similar in AGENTS.md, + but for Phase 1 we will just look for lines that look like paths + in a specific section or just use the defaults if not easily parsable. + Actually, the requirement says "if ResourceLoader finds a different list in AGENTS.md... + it has priority... but never reduce". + Let's implement a simple parser for AGENTS.md for protected files. + """ + effective = PROTECTED_FILES.copy() + content = self.discover_agents_md() + if content: + # Simple heuristic: look for lines starting with - and containing a path + # or specifically in a "Protected Files" section + # For now, let's look for: - `path/to/file` (protected) + # or just any line with 'protected' and a backticked path + matches = re.findall(r"- `([^`]+)`.*protected", content, re.IGNORECASE) + for m in matches: + effective.add(m) + return effective diff --git a/very-simplified-stack/cognito-backend/app/core/tools/__init__.py b/very-simplified-stack/cognito-backend/app/core/tools/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/very-simplified-stack/cognito-backend/app/core/tools/base.py b/very-simplified-stack/cognito-backend/app/core/tools/base.py new file mode 100644 index 0000000..1a4ebb1 --- /dev/null +++ b/very-simplified-stack/cognito-backend/app/core/tools/base.py @@ -0,0 +1,21 @@ +from abc import ABC, abstractmethod +from typing import Any, Set +from pydantic import BaseModel + +class ToolContext(BaseModel): + cwd: str + trusted: bool + protected_files: Set[str] + +class ToolResult(BaseModel): + output: str + is_error: bool = False + +class AgentTool(ABC): + name: str + description: str + parameters_schema: dict[str, Any] # Standard JSON Schema + + @abstractmethod + async def execute(self, arguments: dict[str, Any], context: ToolContext) -> ToolResult: + ... diff --git a/very-simplified-stack/cognito-backend/app/core/tools/bash_tool.py b/very-simplified-stack/cognito-backend/app/core/tools/bash_tool.py new file mode 100644 index 0000000..bd2cb0e --- /dev/null +++ b/very-simplified-stack/cognito-backend/app/core/tools/bash_tool.py @@ -0,0 +1,55 @@ +import asyncio +import os +import re +from typing import Any, Dict +from app.core.tools.base import AgentTool, ToolContext, ToolResult + +class BashTool(AgentTool): + name = "bash" + description = "Execute a bash command in the current workspace." + parameters_schema = { + "type": "object", + "properties": { + "command": {"type": "string", "description": "The bash command to execute."}, + "timeout_seconds": {"type": "integer", "description": "Command timeout in seconds.", "default": 30}, + }, + "required": ["command"], + } + + async def execute(self, arguments: Dict[str, Any], context: ToolContext) -> ToolResult: + command = arguments.get("command") + timeout = min(int(arguments.get("timeout_seconds", 30)), 120) + + if not command: + return ToolResult(is_error=True, output="Error: command is required.") + + # sudo rejection + if re.search(r"\bsudo\b", command, re.IGNORECASE): + return ToolResult(is_error=True, output="Error: Use of 'sudo' is strictly forbidden.") + + # Trust check + if not context.trusted: + return ToolResult(is_error=True, output="Proyecto no confiado (untrusted). Ejecuta project-trust set antes de ejecutar comandos.") + + try: + # Run the command with a timeout + process = await asyncio.create_subprocess_shell( + command, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=context.cwd + ) + + try: + stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout) + output = (stdout + stderr).decode("utf-8") + return ToolResult(output=output) + except asyncio.TimeoutError: + try: + process.kill() + except ProcessLookupError: + pass + return ToolResult(is_error=True, output=f"Error: Command timed out after {timeout} seconds.") + + except Exception as e: + return ToolResult(is_error=True, output=f"Error executing command: {str(e)}") diff --git a/very-simplified-stack/cognito-backend/app/core/tools/edit_tool.py b/very-simplified-stack/cognito-backend/app/core/tools/edit_tool.py new file mode 100644 index 0000000..68420cc --- /dev/null +++ b/very-simplified-stack/cognito-backend/app/core/tools/edit_tool.py @@ -0,0 +1,61 @@ +import os +from typing import Any, Dict +from app.core.tools.base import AgentTool, ToolContext, ToolResult + +class EditTool(AgentTool): + name = "edit" + description = "Edit a file by replacing a specific string with a new one." + parameters_schema = { + "type": "object", + "properties": { + "path": {"type": "string", "description": "Path to the file relative to cwd."}, + "old_str": {"type": "string", "description": "The exact string to be replaced."}, + "new_str": {"type": "string", "description": "The string to replace old_str with."}, + }, + "required": ["path", "old_str", "new_str"], + } + + async def execute(self, arguments: Dict[str, Any], context: ToolContext) -> ToolResult: + path = arguments.get("path") + old_str = arguments.get("old_str") + new_str = arguments.get("new_str") + + if not path or old_str is None or new_str is None: + return ToolResult(is_error=True, output="Error: path, old_str and new_str are required.") + + abs_cwd = os.path.realpath(context.cwd) + norm_path = os.path.normpath(path) + target_path = os.path.realpath(os.path.join(abs_cwd, norm_path)) + + # Path traversal protection + if not target_path.startswith(abs_cwd): + return ToolResult(is_error=True, output=f"Error: Access denied. Path '{path}' is outside of workspace.") + + # Protected files check + if norm_path in context.protected_files: + return ToolResult(is_error=True, output=f"Archivo protegido: {norm_path}. No se puede modificar vía agente.") + + # Trust check + if not context.trusted: + return ToolResult(is_error=True, output="Proyecto no confiado (untrusted). Ejecuta project-trust set antes de escribir.") + + if not os.path.exists(target_path): + return ToolResult(is_error=True, output=f"Error: File '{path}' not found.") + + try: + with open(target_path, "r", encoding="utf-8") as f: + content = f.read() + + count = content.count(old_str) + if count == 0: + return ToolResult(is_error=True, output=f"Error: old_str not found in '{path}'.") + if count > 1: + return ToolResult(is_error=True, output=f"Error: old_str appears {count} times in '{path}'. Must be unique.") + + new_content = content.replace(old_str, new_str) + with open(target_path, "w", encoding="utf-8") as f: + f.write(new_content) + + return ToolResult(output=f"File '{norm_path}' edited successfully.") + except Exception as e: + return ToolResult(is_error=True, output=f"Error editing file: {str(e)}") diff --git a/very-simplified-stack/cognito-backend/app/core/tools/read_tool.py b/very-simplified-stack/cognito-backend/app/core/tools/read_tool.py new file mode 100644 index 0000000..255f9ef --- /dev/null +++ b/very-simplified-stack/cognito-backend/app/core/tools/read_tool.py @@ -0,0 +1,45 @@ +import os +from typing import Any, Dict +from app.core.tools.base import AgentTool, ToolContext, ToolResult + +class ReadTool(AgentTool): + name = "read" + description = "Read the content of a file. Supports optional offset and limit for large files." + parameters_schema = { + "type": "object", + "properties": { + "path": {"type": "string", "description": "Path to the file relative to cwd."}, + "offset": {"type": "integer", "description": "Character offset to start reading from.", "default": 0}, + "limit": {"type": "integer", "description": "Maximum number of characters to read.", "default": 10000}, + }, + "required": ["path"], + } + + async def execute(self, arguments: Dict[str, Any], context: ToolContext) -> ToolResult: + path = arguments.get("path") + offset = arguments.get("offset", 0) + limit = arguments.get("limit", 10000) + + if not path: + return ToolResult(is_error=True, output="Error: path is required.") + + # Path traversal protection + abs_cwd = os.path.realpath(context.cwd) + target_path = os.path.realpath(os.path.join(abs_cwd, path)) + + if not target_path.startswith(abs_cwd): + return ToolResult(is_error=True, output=f"Error: Access denied. Path '{path}' is outside of workspace.") + + if not os.path.exists(target_path): + return ToolResult(is_error=True, output=f"Error: File '{path}' not found.") + + if os.path.isdir(target_path): + return ToolResult(is_error=True, output=f"Error: '{path}' is a directory.") + + try: + with open(target_path, "r", encoding="utf-8") as f: + f.seek(offset) + content = f.read(limit) + return ToolResult(output=content) + except Exception as e: + return ToolResult(is_error=True, output=f"Error reading file: {str(e)}") diff --git a/very-simplified-stack/cognito-backend/app/core/tools/write_tool.py b/very-simplified-stack/cognito-backend/app/core/tools/write_tool.py new file mode 100644 index 0000000..4fd7105 --- /dev/null +++ b/very-simplified-stack/cognito-backend/app/core/tools/write_tool.py @@ -0,0 +1,46 @@ +import os +from typing import Any, Dict +from app.core.tools.base import AgentTool, ToolContext, ToolResult + +class WriteTool(AgentTool): + name = "write" + description = "Write or overwrite a file with the provided content." + parameters_schema = { + "type": "object", + "properties": { + "path": {"type": "string", "description": "Path to the file relative to cwd."}, + "content": {"type": "string", "description": "Content to write to the file."}, + }, + "required": ["path", "content"], + } + + async def execute(self, arguments: Dict[str, Any], context: ToolContext) -> ToolResult: + path = arguments.get("path") + content = arguments.get("content") + + if not path or content is None: + return ToolResult(is_error=True, output="Error: path and content are required.") + + abs_cwd = os.path.realpath(context.cwd) + norm_path = os.path.normpath(path) + target_path = os.path.realpath(os.path.join(abs_cwd, norm_path)) + + # Path traversal protection + if not target_path.startswith(abs_cwd): + return ToolResult(is_error=True, output=f"Error: Access denied. Path '{path}' is outside of workspace.") + + # Protected files check + if norm_path in context.protected_files: + return ToolResult(is_error=True, output=f"Archivo protegido: {norm_path}. No se puede modificar vía agente.") + + # Trust check + if not context.trusted: + return ToolResult(is_error=True, output="Proyecto no confiado (untrusted). Ejecuta project-trust set antes de escribir.") + + try: + os.makedirs(os.path.dirname(target_path), exist_ok=True) + with open(target_path, "w", encoding="utf-8") as f: + f.write(content) + return ToolResult(output=f"File '{norm_path}' written successfully.") + except Exception as e: + return ToolResult(is_error=True, output=f"Error writing file: {str(e)}") diff --git a/very-simplified-stack/cognito-backend/app/core/uncertainty.py b/very-simplified-stack/cognito-backend/app/core/uncertainty.py new file mode 100644 index 0000000..0d357b4 --- /dev/null +++ b/very-simplified-stack/cognito-backend/app/core/uncertainty.py @@ -0,0 +1,39 @@ +import math +import logging +from typing import Any, Optional + +logger = logging.getLogger(__name__) + +def compute_uncertainty(logprob_data: Any) -> Optional[float]: + """ + Shannon entropy of the top-k distribution, normalized to [0, 1]. + Expects Ollama-style logprobs data. + """ + if not logprob_data: + return None + + top_logprobs = {} + if isinstance(logprob_data, list) and logprob_data: + # Ollama returns a list of token info, usually just one for the current token + entry = logprob_data[0] + for candidate in entry.get("top_logprobs", []): + top_logprobs[candidate["token"]] = candidate["logprob"] + elif isinstance(logprob_data, dict): + # Already parsed or different format + top_logprobs = logprob_data + + if not top_logprobs: + return None + + try: + probs = [math.exp(lp) for lp in top_logprobs.values()] + total = sum(probs) + if total == 0: + return 0.0 + probs = [p / total for p in probs] + entropy = -sum(p * math.log2(p) for p in probs if p > 0) + max_entropy = math.log2(len(probs)) if len(probs) > 1 else 1.0 + return entropy / max_entropy if max_entropy > 0 else 0.0 + except Exception as e: + logger.error("[Uncertainty] Error computing entropy: %s", e) + return None diff --git a/very-simplified-stack/cognito-backend/app/services/backend_client.py b/very-simplified-stack/cognito-backend/app/services/backend_client.py index 1259fef..de99536 100644 --- a/very-simplified-stack/cognito-backend/app/services/backend_client.py +++ b/very-simplified-stack/cognito-backend/app/services/backend_client.py @@ -2,7 +2,7 @@ backend_client.py A unified async client that can talk to: - - Ollama native API (POST /api/generate) + - Ollama native API (POST /api/generate or POST /api/chat) - OpenAI-compatible (POST /v1/chat/completions) Supports both blocking and streaming responses. @@ -12,13 +12,16 @@ import httpx import json import logging -from typing import Any, AsyncGenerator, Dict, Optional +from typing import Any, AsyncGenerator, Dict, List, Optional, Set from app.services.backend_registry import BackendConfig, BackendType logger = logging.getLogger(__name__) +# In-memory cache for tool calling support: {(backend_name, model_name): bool} +_TOOL_SUPPORT_CACHE: Dict[tuple[str, str], bool] = {} + class BackendClient: """ Thin async wrapper around a single BackendConfig. @@ -70,6 +73,53 @@ async def generate_stream( async for chunk in self._stream_openai(prompt, model_params): yield chunk + async def generate_with_tools( + self, + messages: List[Dict[str, Any]], + tools: List[Dict[str, Any]], + model_params: Optional[Dict[str, Any]] = None, + ) -> AsyncGenerator[Dict[str, Any], None]: + """ + New method for Phase 1. Calls backend with tool support. + If tools aren't supported, yields an error event or signals failure. + Uses POST /api/chat for Ollama. + """ + cache_key = (self.config.name, self.config.model) + if _TOOL_SUPPORT_CACHE.get(cache_key) is False: + raise NotImplementedError(f"Backend '{self.config.name}' does not support tool calling.") + + try: + if self.config.backend_type == BackendType.OLLAMA: + async for chunk in self._stream_ollama_chat(messages, tools, model_params): + yield chunk + else: + async for chunk in self._stream_openai_chat(messages, tools, model_params): + yield chunk + + # If we reached here without error, mark as supported + if cache_key not in _TOOL_SUPPORT_CACHE: + _TOOL_SUPPORT_CACHE[cache_key] = True + + except (httpx.HTTPStatusError, httpx.RequestError, NotImplementedError) as e: + # Detect if it's a "tool calling not supported" error + is_unsupported = False + if isinstance(e, httpx.HTTPStatusError): + try: + error_data = e.response.json() + error_msg = str(error_data).lower() + if "tool" in error_msg or "schema" in error_msg: + is_unsupported = True + except: + pass + elif isinstance(e, NotImplementedError): + is_unsupported = True + + if is_unsupported: + _TOOL_SUPPORT_CACHE[cache_key] = False + logger.warning("Backend '%s' marked as NO TOOL SUPPORT.", self.config.name) + + raise e + async def check_health(self) -> bool: """ Returns True if the backend is reachable. @@ -236,3 +286,95 @@ async def _stream_openai( if token or logprobs: yield {"token": token, "logprobs": logprobs} + + # ── Tool calling helpers ─────────────────────────────────────────────────── + + async def _stream_ollama_chat( + self, + messages: List[Dict[str, Any]], + tools: List[Dict[str, Any]], + model_params: Optional[Dict[str, Any]], + ) -> AsyncGenerator[Dict[str, Any], None]: + url = f"{self.config.base_url}/api/chat" + payload: Dict[str, Any] = { + "model": self.config.model, + "messages": messages, + "tools": tools, + "stream": True, + } + if model_params: + payload.update({k: v for k, v in model_params.items() if k != "stream"}) + + logger.info("[%s] CHAT STREAM %s (ollama tools)", self.config.name, url) + + async with self._client.stream("POST", url, json=payload) as response: + response.raise_for_status() + async for line in response.aiter_lines(): + if not line.strip(): + continue + try: + data = json.loads(line) + except json.JSONDecodeError: + continue + + msg = data.get("message", {}) + token = msg.get("content", "") + tool_calls = msg.get("tool_calls") + logprobs = data.get("logprobs") + + # Ollama returns tool_calls in a single chunk + yield { + "token": token, + "tool_calls": tool_calls, + "logprobs": logprobs, + "done": data.get("done", False) + } + + if data.get("done", False): + break + + async def _stream_openai_chat( + self, + messages: List[Dict[str, Any]], + tools: List[Dict[str, Any]], + model_params: Optional[Dict[str, Any]], + ) -> AsyncGenerator[Dict[str, Any], None]: + url = f"{self.config.base_url}/v1/chat/completions" + payload: Dict[str, Any] = { + "model": self.config.model, + "messages": messages, + "tools": tools, + "stream": True, + } + if model_params: + payload.update({k: v for k, v in model_params.items() if k != "stream"}) + + logger.info("[%s] CHAT STREAM %s (openai tools)", self.config.name, url) + + async with self._client.stream("POST", url, json=payload) as response: + response.raise_for_status() + async for line in response.aiter_lines(): + if not line.strip(): + continue + if line.startswith("data:"): + line = line[5:].strip() + if line == "[DONE]": + break + try: + data = json.loads(line) + except json.JSONDecodeError: + continue + + try: + choice = data["choices"][0] + delta = choice.get("delta", {}) + token = delta.get("content", "") + tool_calls = delta.get("tool_calls") + logprobs = choice.get("logprobs") + except (KeyError, IndexError): + token = "" + tool_calls = None + logprobs = None + + if token or tool_calls or logprobs: + yield {"token": token, "tool_calls": tool_calls, "logprobs": logprobs} diff --git a/very-simplified-stack/cognito-backend/app/services/backend_router.py b/very-simplified-stack/cognito-backend/app/services/backend_router.py index 8f99dc9..be2cd5a 100644 --- a/very-simplified-stack/cognito-backend/app/services/backend_router.py +++ b/very-simplified-stack/cognito-backend/app/services/backend_router.py @@ -18,7 +18,7 @@ import asyncio import logging -from typing import Any, Dict, Optional +from typing import Any, AsyncGenerator, Dict, List, Optional from app.services.backend_client import BackendClient from app.services.backend_registry import BACKENDS_BY_PRIORITY, BackendConfig @@ -79,6 +79,52 @@ async def generate( f"Last error: {last_exc}" ) from last_exc + async def generate_with_tools( + self, + messages: List[Dict[str, Any]], + tools: List[Dict[str, Any]], + model_params: Optional[Dict[str, Any]] = None, + ) -> AsyncGenerator[Dict[str, Any], None]: + """ + Try each backend in priority order for tool-enabled chat. + If a backend doesn't support tools, it's skipped. + """ + last_exc: Optional[Exception] = None + + for attempt, client in enumerate(self._clients, start=1): + try: + logger.info( + "[Tools] Attempt %d/%d → backend '%s'", + attempt, len(self._clients), client.config.name, + ) + async for chunk in client.generate_with_tools(messages, tools, model_params): + yield chunk + return # Success! + + except (NotImplementedError, RuntimeError) as exc: + # Handle NotImplementedError (explicitly marked unsupported) + # or general errors that we might want to failover from + logger.warning( + "[Tools] Backend '%s' skipped (attempt %d): %s", + client.config.name, attempt, exc, + ) + last_exc = exc + continue + except Exception as exc: + logger.warning( + "[Tools] Backend '%s' failed (attempt %d): %s", + client.config.name, attempt, exc, + ) + last_exc = exc + continue + + # All backends exhausted + logger.error("[Tools] All %d backends failed or unsupported.", len(self._clients)) + raise RuntimeError( + f"All backends failed or do not support tool calling. " + f"Last error: {last_exc}" + ) + async def health_all(self) -> Dict[str, str]: """ Concurrently health-checks every backend. diff --git a/very-simplified-stack/cognito-backend/tests/test_agent_loop.py b/very-simplified-stack/cognito-backend/tests/test_agent_loop.py new file mode 100644 index 0000000..4a9b439 --- /dev/null +++ b/very-simplified-stack/cognito-backend/tests/test_agent_loop.py @@ -0,0 +1,90 @@ +import pytest +import json +from unittest.mock import AsyncMock, MagicMock +from app.core.agent_loop import agent_loop +from app.core.tools.base import ToolContext, ToolResult, AgentTool +from app.core.events import TextDeltaEvent, ToolCallEvent, ToolResultEvent, DoneEvent + +class MockTool(AgentTool): + name = "mock_tool" + description = "A mock tool" + parameters_schema = {"type": "object", "properties": {"arg1": {"type": "string"}}} + + def __init__(self): + self.mock_execute = AsyncMock() + + async def execute(self, arguments, context): + return await self.mock_execute(arguments, context) + +@pytest.fixture +def tool_context(): + return ToolContext(cwd="/tmp", trusted=True, protected_files=set()) + +@pytest.mark.asyncio +async def test_agent_loop_simple_text(tool_context): + backend_router = MagicMock() + + async def mock_generate(*args, **kwargs): + yield {"token": "Hello"} + yield {"token": " world"} + + backend_router.generate_with_tools = mock_generate + + messages = [{"role": "user", "content": "Hi"}] + tools = [] + + events = [] + async for event in agent_loop(messages, tools, tool_context, backend_router): + events.append(event) + + assert len(events) == 3 # "Hello", " world", DoneEvent + assert isinstance(events[0], TextDeltaEvent) + assert events[0].content == "Hello" + assert isinstance(events[1], TextDeltaEvent) + assert events[1].content == " world" + assert isinstance(events[2], DoneEvent) + assert events[2].stop_reason == "end_turn" + +@pytest.mark.asyncio +async def test_agent_loop_tool_call(tool_context): + backend_router = MagicMock() + mock_tool = MockTool() + mock_tool.mock_execute.return_value = ToolResult(output="tool result") + + # Turn 1: model calls tool + # Turn 2: model responds with final text + turn = 0 + async def mock_generate(*args, **kwargs): + nonlocal turn + turn += 1 + if turn == 1: + yield { + "token": "I will use a tool", + "tool_calls": [{ + "function": {"name": "mock_tool", "arguments": {"arg1": "val1"}} + }] + } + else: + yield {"token": "Final answer"} + + backend_router.generate_with_tools = mock_generate + + messages = [{"role": "user", "content": "use tool"}] + tools = [mock_tool] + + events = [] + async for event in agent_loop(messages, tools, tool_context, backend_router): + events.append(event) + + # Events expected: + # 1. TextDeltaEvent "I will use a tool" + # 2. ToolCallEvent "mock_tool" + # 3. ToolResultEvent "tool result" + # 4. TextDeltaEvent "Final answer" + # 5. DoneEvent "end_turn" + + assert any(isinstance(e, ToolCallEvent) for e in events) + assert any(isinstance(e, ToolResultEvent) for e in events) + assert any(isinstance(e, TextDeltaEvent) and e.content == "Final answer" for e in events) + assert events[-1].stop_reason == "end_turn" + mock_tool.mock_execute.assert_called_once() diff --git a/very-simplified-stack/cognito-backend/tests/test_project_trust.py b/very-simplified-stack/cognito-backend/tests/test_project_trust.py new file mode 100644 index 0000000..52cc7de --- /dev/null +++ b/very-simplified-stack/cognito-backend/tests/test_project_trust.py @@ -0,0 +1,17 @@ +import os +import tempfile +from pathlib import Path +from app.core.project_trust import ProjectTrustStore + +def test_project_trust_store(): + with tempfile.NamedTemporaryFile() as tmp: + store = ProjectTrustStore(store_path=Path(tmp.name)) + repo = "/tmp/fake-repo" + + assert store.is_trusted(repo) is False + + store.set_trusted(repo, True) + assert store.is_trusted(repo) is True + + store.set_trusted(repo, False) + assert store.is_trusted(repo) is False diff --git a/very-simplified-stack/cognito-backend/tests/test_tools.py b/very-simplified-stack/cognito-backend/tests/test_tools.py new file mode 100644 index 0000000..1b7b003 --- /dev/null +++ b/very-simplified-stack/cognito-backend/tests/test_tools.py @@ -0,0 +1,95 @@ +import pytest +import os +import tempfile +from pathlib import Path +from app.core.project_trust import ProjectTrustStore +from app.core.tools.base import ToolContext, ToolResult +from app.core.tools.read_tool import ReadTool +from app.core.tools.write_tool import WriteTool +from app.core.tools.edit_tool import EditTool +from app.core.tools.bash_tool import BashTool + +@pytest.fixture +def temp_workspace(): + with tempfile.TemporaryDirectory() as tmpdir: + yield tmpdir + +@pytest.fixture +def tool_context(temp_workspace): + return ToolContext( + cwd=temp_workspace, + trusted=True, + protected_files={"protected.txt"} + ) + +@pytest.mark.asyncio +async def test_read_tool(temp_workspace, tool_context): + test_file = Path(temp_workspace) / "test.txt" + test_file.write_text("hello world") + + tool = ReadTool() + # Success + result = await tool.execute({"path": "test.txt"}, tool_context) + assert result.output == "hello world" + assert not result.is_error + + # Path traversal + result = await tool.execute({"path": "../test.txt"}, tool_context) + assert result.is_error + assert "outside of workspace" in result.output + +@pytest.mark.asyncio +async def test_write_tool(temp_workspace, tool_context): + tool = WriteTool() + + # Success + result = await tool.execute({"path": "new.txt", "content": "data"}, tool_context) + assert not result.is_error + assert (Path(temp_workspace) / "new.txt").read_text() == "data" + + # Protected file + result = await tool.execute({"path": "protected.txt", "content": "data"}, tool_context) + assert result.is_error + assert "Archivo protegido" in result.output + + # Untrusted + tool_context.trusted = False + result = await tool.execute({"path": "untrusted.txt", "content": "data"}, tool_context) + assert result.is_error + assert "no confiado" in result.output + +@pytest.mark.asyncio +async def test_edit_tool(temp_workspace, tool_context): + test_file = Path(temp_workspace) / "edit.txt" + test_file.write_text("line 1\nline 2\nline 3") + + tool = EditTool() + # Success + result = await tool.execute({"path": "edit.txt", "old_str": "line 2", "new_str": "replaced"}, tool_context) + assert not result.is_error + assert "replaced" in test_file.read_text() + + # Not unique + test_file.write_text("aaa\naaa") + result = await tool.execute({"path": "edit.txt", "old_str": "aaa", "new_str": "bbb"}, tool_context) + assert result.is_error + assert "appears 2 times" in result.output + +@pytest.mark.asyncio +async def test_bash_tool(temp_workspace, tool_context): + tool = BashTool() + + # Success + result = await tool.execute({"command": "echo 'hello'"}, tool_context) + assert not result.is_error + assert "hello" in result.output + + # sudo rejection + result = await tool.execute({"command": "sudo ls"}, tool_context) + assert result.is_error + assert "strictly forbidden" in result.output + + # Timeout (mocking or using sleep) + result = await tool.execute({"command": "sleep 2", "timeout_seconds": 1}, tool_context) + assert result.is_error + assert "timed out" in result.output