From f64bcbe89c0d913901b7d1f012df471861051963 Mon Sep 17 00:00:00 2001 From: Claude Code Date: Tue, 25 Aug 2026 08:29:39 +0000 Subject: [PATCH 1/3] chore(fuze): sync vendored repo-manifest schema with canonical Vendored .fuze/repo-manifest.schema.json was stale relative to the canonical copy at izzywdev/FuzeSDLC@main:governance/repo-manifest.schema.json. Replacing with the canonical bytes so gate-manifest (when enabled) does not reject fields the canonical schema allows. Co-Authored-By: Claude Opus 5 --- .fuze/repo-manifest.schema.json | 40 ++++++++++++++++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/.fuze/repo-manifest.schema.json b/.fuze/repo-manifest.schema.json index 3db2bbe..38496c7 100644 --- a/.fuze/repo-manifest.schema.json +++ b/.fuze/repo-manifest.schema.json @@ -231,8 +231,12 @@ "platformAuth": { "type": "object", "additionalProperties": false, - "description": "NEW BLOCK \u2014 no repo declares it yet, and that is the point: it gates the platform-auth capability. Consume @fuzefront/auth (published as @izzywdev/fuzefront-auth) rather than a bespoke verifier. A product NEVER calls Permit directly; it knows exactly one thing, the base URL of FuzeFront's Security API.", + "description": "NEW BLOCK. Consume @fuzefront/auth (published as @izzywdev/fuzefront-auth) rather than a bespoke verifier. A product NEVER calls Permit directly; it knows exactly one thing, the base URL of FuzeFront's Security API. gate-platform-auth ENFORCES BY DEFAULT \u2014 this block is how a repo opts OUT, not how it opts in.", "properties": { + "enforce": { + "type": "boolean", + "description": "Ratchet for gate-platform-auth, and it is OPT-OUT: absent means ENFORCING. Set false only to silence the gate while a repo migrates, and only together with `reason` \u2014 an `enforce: false` with no reason is ignored and the gate enforces anyway, because an undocumented opt-out is indistinguishable from an oversight. The earlier opt-in shape was chosen to avoid redding the fleet on pre-existing violations, but that is how gate-identifier reached zero adoption across 21 repos: a check nobody enabled is indistinguishable from a check that does not exist. What actually prevents a `|| true` is visibility, not coldness \u2014 an `enforce: false` naming a repo and a reason is greppable and countable; `|| true` in a workflow is neither." + }, "mode": { "enum": [ "federated-jwks", @@ -251,6 +255,10 @@ }, "note": { "type": "string" + }, + "reason": { + "type": "string", + "description": "REQUIRED when enforce is false. What blocks adoption and who owns closing it. This is the whole cost of the escape hatch: the opt-out must read as debt someone wrote down, not as a setting someone left alone." } } }, @@ -667,6 +675,36 @@ } } }, + "dataTier": { + "type": "array", + "description": "Declarative data-tier provisioning request (the IaC hand-off to FuzeInfra). FuzeInfra's reconciler consumes each entry: it ensures the per-service role exists AND is GRANTED the declared privileges on the declared database, then VERIFIES the role can actually read/write it (fail-loud if a role can auth but not access its DB). Replaces the old ad-hoc '@claude please provision' request (governance/shared-cluster-deploy.md §5). Every store the product's role authenticates to MUST be declared here, with the exact database name the app uses — a role granted on the wrong db name is the classic silent-empty-data bug.", + "items": { + "type": "object", + "additionalProperties": false, + "required": ["store", "database", "role"], + "properties": { + "store": { "enum": ["postgres", "mongo", "redis", "neo4j", "chroma"], "description": "Shared datastore this role needs access in." }, + "database": { "type": "string", "description": "The exact database/keyspace name the app reads/writes (e.g. robot_catalog). The role MUST be granted on THIS name; provisioning verifies it." }, + "role": { "type": "string", "description": "The per-service role/user (e.g. mendys)." }, + "privileges": { "enum": ["readWrite", "read", "admin"], "default": "readWrite", "description": "Privilege level to grant the role on `database`." }, + "authSource": { "type": "string", "description": "Mongo authSource db the role authenticates against (e.g. admin), when it differs from `database`." } + } + } + }, + "egress": { + "type": "array", + "description": "External hosts the product's pods need outbound HTTPS to. The shared cluster is egress-restricted (HTTP-only behind the Cloudflare tunnel; no default outbound to third-party APIs), so every external dependency MUST be declared here. FuzeInfra's reconciler turns these into namespace egress allow-rules (NetworkPolicy / egress gateway). Declare each third-party API explicitly (e.g. LLM providers).", + "items": { + "type": "object", + "additionalProperties": false, + "required": ["host"], + "properties": { + "host": { "type": "string", "description": "FQDN, e.g. api.openai.com." }, + "port": { "type": "integer", "default": 443, "description": "Destination port (default 443)." }, + "reason": { "type": "string", "description": "Why the product needs it (e.g. 'AI keyword generation')." } + } + } + }, "dependsOn": { "type": "array", "description": "Product-to-product dependencies this repo consumes beyond the spine (e.g. FuzeService dependsOn FuzeContact, FuzeBI).", From b09b32a6897c37703c2ffb8fc95b84deb9db6c9a Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Tue, 25 Aug 2026 08:35:14 +0000 Subject: [PATCH 2/3] chore(governance): reconcile managed files to FuzeSDLC v1 [skip ci] --- .fuze/repo-manifest.schema.json | 40 +- hierarchy_endpoints.py | 862 +- .../orchestrator/agent_expertise_tracker.py | 1040 +- services/orchestrator/claude_code_wrapper.py | 1684 +-- services/orchestrator/claude_sdk_manager.py | 1004 +- .../context_enhancement_service.py | 1548 +- services/orchestrator/context_service.py | 304 +- services/orchestrator/conversation_manager.py | 1204 +- .../orchestrator/coordination_endpoints.py | 1250 +- .../orchestrator/goal_conversation_service.py | 2080 +-- services/orchestrator/hierarchy_endpoints.py | 786 +- .../knowledge_propagation_engine.py | 1916 +-- services/orchestrator/main.py | 12026 ++++++++-------- services/orchestrator/mcp_integration.py | 1318 +- services/orchestrator/model_configuration.py | 1116 +- .../orchestrator/multi_agent_coordinator.py | 1876 +-- .../orchestrator/task_execution_engine.py | 2612 ++-- .../orchestrator/task_knowledge_extractor.py | 1758 +-- services/orchestrator/task_queue.py | 248 +- .../orchestrator/team_knowledge_manager.py | 1738 +-- 20 files changed, 18186 insertions(+), 18224 deletions(-) diff --git a/.fuze/repo-manifest.schema.json b/.fuze/repo-manifest.schema.json index 38496c7..3db2bbe 100644 --- a/.fuze/repo-manifest.schema.json +++ b/.fuze/repo-manifest.schema.json @@ -231,12 +231,8 @@ "platformAuth": { "type": "object", "additionalProperties": false, - "description": "NEW BLOCK. Consume @fuzefront/auth (published as @izzywdev/fuzefront-auth) rather than a bespoke verifier. A product NEVER calls Permit directly; it knows exactly one thing, the base URL of FuzeFront's Security API. gate-platform-auth ENFORCES BY DEFAULT \u2014 this block is how a repo opts OUT, not how it opts in.", + "description": "NEW BLOCK \u2014 no repo declares it yet, and that is the point: it gates the platform-auth capability. Consume @fuzefront/auth (published as @izzywdev/fuzefront-auth) rather than a bespoke verifier. A product NEVER calls Permit directly; it knows exactly one thing, the base URL of FuzeFront's Security API.", "properties": { - "enforce": { - "type": "boolean", - "description": "Ratchet for gate-platform-auth, and it is OPT-OUT: absent means ENFORCING. Set false only to silence the gate while a repo migrates, and only together with `reason` \u2014 an `enforce: false` with no reason is ignored and the gate enforces anyway, because an undocumented opt-out is indistinguishable from an oversight. The earlier opt-in shape was chosen to avoid redding the fleet on pre-existing violations, but that is how gate-identifier reached zero adoption across 21 repos: a check nobody enabled is indistinguishable from a check that does not exist. What actually prevents a `|| true` is visibility, not coldness \u2014 an `enforce: false` naming a repo and a reason is greppable and countable; `|| true` in a workflow is neither." - }, "mode": { "enum": [ "federated-jwks", @@ -255,10 +251,6 @@ }, "note": { "type": "string" - }, - "reason": { - "type": "string", - "description": "REQUIRED when enforce is false. What blocks adoption and who owns closing it. This is the whole cost of the escape hatch: the opt-out must read as debt someone wrote down, not as a setting someone left alone." } } }, @@ -675,36 +667,6 @@ } } }, - "dataTier": { - "type": "array", - "description": "Declarative data-tier provisioning request (the IaC hand-off to FuzeInfra). FuzeInfra's reconciler consumes each entry: it ensures the per-service role exists AND is GRANTED the declared privileges on the declared database, then VERIFIES the role can actually read/write it (fail-loud if a role can auth but not access its DB). Replaces the old ad-hoc '@claude please provision' request (governance/shared-cluster-deploy.md §5). Every store the product's role authenticates to MUST be declared here, with the exact database name the app uses — a role granted on the wrong db name is the classic silent-empty-data bug.", - "items": { - "type": "object", - "additionalProperties": false, - "required": ["store", "database", "role"], - "properties": { - "store": { "enum": ["postgres", "mongo", "redis", "neo4j", "chroma"], "description": "Shared datastore this role needs access in." }, - "database": { "type": "string", "description": "The exact database/keyspace name the app reads/writes (e.g. robot_catalog). The role MUST be granted on THIS name; provisioning verifies it." }, - "role": { "type": "string", "description": "The per-service role/user (e.g. mendys)." }, - "privileges": { "enum": ["readWrite", "read", "admin"], "default": "readWrite", "description": "Privilege level to grant the role on `database`." }, - "authSource": { "type": "string", "description": "Mongo authSource db the role authenticates against (e.g. admin), when it differs from `database`." } - } - } - }, - "egress": { - "type": "array", - "description": "External hosts the product's pods need outbound HTTPS to. The shared cluster is egress-restricted (HTTP-only behind the Cloudflare tunnel; no default outbound to third-party APIs), so every external dependency MUST be declared here. FuzeInfra's reconciler turns these into namespace egress allow-rules (NetworkPolicy / egress gateway). Declare each third-party API explicitly (e.g. LLM providers).", - "items": { - "type": "object", - "additionalProperties": false, - "required": ["host"], - "properties": { - "host": { "type": "string", "description": "FQDN, e.g. api.openai.com." }, - "port": { "type": "integer", "default": 443, "description": "Destination port (default 443)." }, - "reason": { "type": "string", "description": "Why the product needs it (e.g. 'AI keyword generation')." } - } - } - }, "dependsOn": { "type": "array", "description": "Product-to-product dependencies this repo consumes beyond the spine (e.g. FuzeService dependsOn FuzeContact, FuzeBI).", diff --git a/hierarchy_endpoints.py b/hierarchy_endpoints.py index acaa613..e48d494 100644 --- a/hierarchy_endpoints.py +++ b/hierarchy_endpoints.py @@ -1,432 +1,432 @@ -#!/usr/bin/env python3 -""" -Simple FastAPI service that adds organization/team endpoints -and proxies other requests to the orchestrator -""" - -import asyncio -import asyncpg -import json -import uuid -import httpx -from contextlib import asynccontextmanager -from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect, Depends -from fastapi.middleware.cors import CORSMiddleware -from fastapi.responses import JSONResponse -from pydantic import BaseModel -from typing import List, Optional, Dict, Set - -# Configuration -import os -# SECURITY (issue #6): do not ship a real-looking DB password as a default. -DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://postgres:postgres@localhost:5434/ai_context") -ORCHESTRATOR_URL = os.getenv("ORCHESTRATOR_URL", "http://localhost:8000") - -try: - from auth import get_current_user, require_user, require_org_access, CurrentUser, authenticate_websocket -except Exception: # pragma: no cover - allow import from repo root or service dir - from services.orchestrator.auth import ( # type: ignore - get_current_user, require_user, require_org_access, CurrentUser, - authenticate_websocket, - ) - -@asynccontextmanager -async def lifespan(app: FastAPI): - """Application lifespan. - - Replaces the removed ``app.add_event_handler("startup"/"shutdown", ...)`` - API (dropped in Starlette 1.x). Behaviour is unchanged: open the asyncpg - pool on startup, close it on shutdown. ``startup``/``shutdown`` are - resolved at call time, so they may be defined further down the module. - """ - await startup() - try: - yield - finally: - await shutdown() - - -# SECURITY (issue #6 CRITICAL-1): authenticate every route by default (health -# and docs are on the allowlist inside get_current_user). -app = FastAPI( - title="FuzeAgent Hierarchy API", - version="1.0.0", - dependencies=[Depends(get_current_user)], - lifespan=lifespan, -) - -# SECURITY (issue #6 MEDIUM-2): explicit, non-wildcard origins when credentials -# are allowed (wildcard + credentials is both insecure and spec-invalid). -_cors_origins = [ - o.strip() - for o in os.getenv( - "CORS_ALLOW_ORIGINS", - "http://localhost:3000,http://localhost:3031,http://localhost", - ).split(",") - if o.strip() -] -app.add_middleware( - CORSMiddleware, - allow_origins=_cors_origins, - allow_credentials=True, - allow_methods=["*"], - allow_headers=["*"], -) - -# Database connection pool -db_pool = None - -# WebSocket connection manager -class ConnectionManager: - def __init__(self): - self.active_connections: Set[WebSocket] = set() - - async def connect(self, websocket: WebSocket): - await websocket.accept() - self.active_connections.add(websocket) - - def disconnect(self, websocket: WebSocket): - self.active_connections.discard(websocket) - - async def broadcast(self, message: dict): - disconnected = set() - for connection in self.active_connections: - try: - await connection.send_text(json.dumps(message)) - except: - disconnected.add(connection) - - # Remove disconnected clients - for connection in disconnected: - self.disconnect(connection) - -manager = ConnectionManager() - -async def startup(): - global db_pool - db_pool = await asyncpg.create_pool(DATABASE_URL, min_size=1, max_size=10) - -async def shutdown(): - if db_pool: - await db_pool.close() - -# NOTE: startup/shutdown are wired via the `lifespan` context manager defined -# above and passed to FastAPI(...); `add_event_handler` was removed in -# Starlette 1.x. - -# Pydantic models -class Organization(BaseModel): - id: str - name: str - description: Optional[str] = None - settings: dict = {} - created_at: str - updated_at: str - -class OrganizationCreate(BaseModel): - name: str - description: Optional[str] = None - settings: dict = {} - -class Team(BaseModel): - id: str - organization_id: str - name: str - description: Optional[str] = None - team_type: str = "general" - settings: dict = {} - created_at: str - updated_at: str - -class TeamCreate(BaseModel): - organization_id: str - name: str - description: Optional[str] = None - team_type: str = "general" - settings: dict = {} - -# Organization endpoints -@app.get("/organizations", response_model=List[Organization]) -async def get_organizations(): - async with db_pool.acquire() as conn: - rows = await conn.fetch(""" - SELECT - id::text, name, description, settings, - created_at::text, updated_at::text - FROM organizations - ORDER BY created_at DESC - """) - - organizations = [] - for row in rows: - organizations.append(Organization( - id=row['id'], - name=row['name'], - description=row['description'], - settings=json.loads(row['settings']) if row['settings'] else {}, - created_at=row['created_at'], - updated_at=row['updated_at'] - )) - - return organizations - -@app.post("/organizations", response_model=Organization) -async def create_organization(org_data: OrganizationCreate): - async with db_pool.acquire() as conn: - org_id = str(uuid.uuid4()) - row = await conn.fetchrow(""" - INSERT INTO organizations (id, name, description, settings) - VALUES ($1, $2, $3, $4) - RETURNING - id::text, name, description, settings, - created_at::text, updated_at::text - """, org_id, org_data.name, org_data.description, json.dumps(org_data.settings)) - - organization = Organization( - id=row['id'], - name=row['name'], - description=row['description'], - settings=json.loads(row['settings']) if row['settings'] else {}, - created_at=row['created_at'], - updated_at=row['updated_at'] - ) - - # Broadcast the change - await manager.broadcast({ - "type": "organization_created", - "data": organization.dict() - }) - - return organization - -@app.get("/organizations/{organization_id}", response_model=Organization) -async def get_organization( - organization_id: str, - user: CurrentUser = Depends(require_user), -): - # SECURITY (issue #6 HIGH-2 / BOLA): authorize the specific org id from the - # path; bare ``WHERE id = $1`` is not an authorization boundary. - require_org_access(organization_id, user) - async with db_pool.acquire() as conn: - row = await conn.fetchrow(""" - SELECT - id::text, name, description, settings, - created_at::text, updated_at::text - FROM organizations - WHERE id = $1 - """, organization_id) - - if not row: - raise HTTPException(status_code=404, detail="Organization not found") - - return Organization( - id=row['id'], - name=row['name'], - description=row['description'], - settings=json.loads(row['settings']) if row['settings'] else {}, - created_at=row['created_at'], - updated_at=row['updated_at'] - ) - -# Team endpoints -@app.get("/teams", response_model=List[Team]) -async def get_teams(organization_id: Optional[str] = None): - async with db_pool.acquire() as conn: - if organization_id: - rows = await conn.fetch(""" - SELECT - id::text, organization_id::text, name, description, - team_type, settings, created_at::text, updated_at::text - FROM teams - WHERE organization_id = $1 - ORDER BY created_at DESC - """, organization_id) - else: - rows = await conn.fetch(""" - SELECT - id::text, organization_id::text, name, description, - team_type, settings, created_at::text, updated_at::text - FROM teams - ORDER BY created_at DESC - """) - - teams = [] - for row in rows: - teams.append(Team( - id=row['id'], - organization_id=row['organization_id'], - name=row['name'], - description=row['description'], - team_type=row['team_type'], - settings=json.loads(row['settings']) if row['settings'] else {}, - created_at=row['created_at'], - updated_at=row['updated_at'] - )) - - return teams - -@app.post("/teams", response_model=Team) -async def create_team(team_data: TeamCreate): - async with db_pool.acquire() as conn: - # Verify organization exists - org_exists = await conn.fetchval( - "SELECT EXISTS(SELECT 1 FROM organizations WHERE id = $1)", - team_data.organization_id - ) - if not org_exists: - raise HTTPException(status_code=404, detail="Organization not found") - - team_id = str(uuid.uuid4()) - row = await conn.fetchrow(""" - INSERT INTO teams (id, organization_id, name, description, team_type, settings) - VALUES ($1, $2, $3, $4, $5, $6) - RETURNING - id::text, organization_id::text, name, description, - team_type, settings, created_at::text, updated_at::text - """, team_id, team_data.organization_id, team_data.name, - team_data.description, team_data.team_type, json.dumps(team_data.settings)) - - team = Team( - id=row['id'], - organization_id=row['organization_id'], - name=row['name'], - description=row['description'], - team_type=row['team_type'], - settings=json.loads(row['settings']) if row['settings'] else {}, - created_at=row['created_at'], - updated_at=row['updated_at'] - ) - - # Broadcast the change - await manager.broadcast({ - "type": "team_created", - "data": team.dict() - }) - - return team - -@app.get("/teams/{team_id}", response_model=Team) -async def get_team( - team_id: str, - user: CurrentUser = Depends(require_user), -): - async with db_pool.acquire() as conn: - row = await conn.fetchrow(""" - SELECT - id::text, organization_id::text, name, description, - team_type, settings, created_at::text, updated_at::text - FROM teams - WHERE id = $1 - """, team_id) - - if not row: - raise HTTPException(status_code=404, detail="Team not found") - # SECURITY (issue #6 HIGH-2): authorize via the team's parent org. - require_org_access(row['organization_id'], user) - return Team( - id=row['id'], - organization_id=row['organization_id'], - name=row['name'], - description=row['description'], - team_type=row['team_type'], - settings=json.loads(row['settings']) if row['settings'] else {}, - created_at=row['created_at'], - updated_at=row['updated_at'] - ) - -# WebSocket endpoint for real-time updates -@app.websocket("/ws") -async def websocket_endpoint(websocket: WebSocket): - # SECURITY (issue #6 CRITICAL-2): authenticate BEFORE accept(). The - # app-wide ``dependencies=[Depends(get_current_user)]`` is a no-op on - # WebSocket routes (WS handshakes have no HTTP response channel); every WS - # handler must call authenticate_websocket() first — matching the pattern - # already used by all handlers in services/orchestrator/main.py. - user = await authenticate_websocket(websocket) - if user is None: - return # authenticate_websocket already closed the socket (1008) - await manager.connect(websocket) - try: - while True: - # Keep the connection alive and listen for client pings - data = await websocket.receive_text() - if data == "ping": - await websocket.send_text("pong") - except WebSocketDisconnect: - manager.disconnect(websocket) - -# Proxy all other requests to the original orchestrator -@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"]) -async def proxy_to_orchestrator(path: str, request: Request): - url = f"{ORCHESTRATOR_URL}/{path}" - - # Handle CORS preflight requests - if request.method == "OPTIONS": - return JSONResponse( - content={}, - headers={ - "Access-Control-Allow-Origin": "*", - "Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, PATCH, OPTIONS", - "Access-Control-Allow-Headers": "*", - } - ) - - try: - async with httpx.AsyncClient(timeout=30.0) as client: - # Get request body if present - body = None - if request.method in ["POST", "PUT", "PATCH"]: - body = await request.body() - # Log the request for debugging - if path == "agents/from-template": - print(f"DEBUG: Proxying agents/from-template request") - print(f"DEBUG: URL: {url}") - print(f"DEBUG: Body: {body.decode() if body else 'None'}") - print(f"DEBUG: Headers: {dict(request.headers)}") - - # Forward the request - response = await client.request( - method=request.method, - url=url, - params=request.query_params, - content=body, - headers={k: v for k, v in request.headers.items() - if k.lower() not in ['host', 'content-length']}, - ) - - # Check if this was an agent creation and broadcast the change - if (request.method == "POST" and - (path == "agents" or path == "agents/from-template") and - response.status_code in [200, 201] and - response.headers.get("content-type", "").startswith("application/json")): - try: - agent_data = response.json() - await manager.broadcast({ - "type": "agent_created", - "data": agent_data - }) - except: - pass # Ignore broadcast errors - - return JSONResponse( - content=response.json() if response.headers.get("content-type", "").startswith("application/json") else response.text, - status_code=response.status_code, - headers={ - "Access-Control-Allow-Origin": "*", - **{k: v for k, v in response.headers.items() - if k.lower() not in ['content-length', 'transfer-encoding', 'connection']} - } - ) - - except httpx.RequestError as e: - print(f"ERROR: RequestError in proxy: {str(e)}") - raise HTTPException(status_code=503, detail=f"Failed to connect to orchestrator: {str(e)}") - except Exception as e: - print(f"ERROR: Exception in proxy: {str(e)}") - import traceback - traceback.print_exc() - raise HTTPException(status_code=500, detail=f"Proxy error: {str(e)}") - -if __name__ == "__main__": - import uvicorn +#!/usr/bin/env python3 +""" +Simple FastAPI service that adds organization/team endpoints +and proxies other requests to the orchestrator +""" + +import asyncio +import asyncpg +import json +import uuid +import httpx +from contextlib import asynccontextmanager +from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect, Depends +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import JSONResponse +from pydantic import BaseModel +from typing import List, Optional, Dict, Set + +# Configuration +import os +# SECURITY (issue #6): do not ship a real-looking DB password as a default. +DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://postgres:postgres@localhost:5434/ai_context") +ORCHESTRATOR_URL = os.getenv("ORCHESTRATOR_URL", "http://localhost:8000") + +try: + from auth import get_current_user, require_user, require_org_access, CurrentUser, authenticate_websocket +except Exception: # pragma: no cover - allow import from repo root or service dir + from services.orchestrator.auth import ( # type: ignore + get_current_user, require_user, require_org_access, CurrentUser, + authenticate_websocket, + ) + +@asynccontextmanager +async def lifespan(app: FastAPI): + """Application lifespan. + + Replaces the removed ``app.add_event_handler("startup"/"shutdown", ...)`` + API (dropped in Starlette 1.x). Behaviour is unchanged: open the asyncpg + pool on startup, close it on shutdown. ``startup``/``shutdown`` are + resolved at call time, so they may be defined further down the module. + """ + await startup() + try: + yield + finally: + await shutdown() + + +# SECURITY (issue #6 CRITICAL-1): authenticate every route by default (health +# and docs are on the allowlist inside get_current_user). +app = FastAPI( + title="FuzeAgent Hierarchy API", + version="1.0.0", + dependencies=[Depends(get_current_user)], + lifespan=lifespan, +) + +# SECURITY (issue #6 MEDIUM-2): explicit, non-wildcard origins when credentials +# are allowed (wildcard + credentials is both insecure and spec-invalid). +_cors_origins = [ + o.strip() + for o in os.getenv( + "CORS_ALLOW_ORIGINS", + "http://localhost:3000,http://localhost:3031,http://localhost", + ).split(",") + if o.strip() +] +app.add_middleware( + CORSMiddleware, + allow_origins=_cors_origins, + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +# Database connection pool +db_pool = None + +# WebSocket connection manager +class ConnectionManager: + def __init__(self): + self.active_connections: Set[WebSocket] = set() + + async def connect(self, websocket: WebSocket): + await websocket.accept() + self.active_connections.add(websocket) + + def disconnect(self, websocket: WebSocket): + self.active_connections.discard(websocket) + + async def broadcast(self, message: dict): + disconnected = set() + for connection in self.active_connections: + try: + await connection.send_text(json.dumps(message)) + except: + disconnected.add(connection) + + # Remove disconnected clients + for connection in disconnected: + self.disconnect(connection) + +manager = ConnectionManager() + +async def startup(): + global db_pool + db_pool = await asyncpg.create_pool(DATABASE_URL, min_size=1, max_size=10) + +async def shutdown(): + if db_pool: + await db_pool.close() + +# NOTE: startup/shutdown are wired via the `lifespan` context manager defined +# above and passed to FastAPI(...); `add_event_handler` was removed in +# Starlette 1.x. + +# Pydantic models +class Organization(BaseModel): + id: str + name: str + description: Optional[str] = None + settings: dict = {} + created_at: str + updated_at: str + +class OrganizationCreate(BaseModel): + name: str + description: Optional[str] = None + settings: dict = {} + +class Team(BaseModel): + id: str + organization_id: str + name: str + description: Optional[str] = None + team_type: str = "general" + settings: dict = {} + created_at: str + updated_at: str + +class TeamCreate(BaseModel): + organization_id: str + name: str + description: Optional[str] = None + team_type: str = "general" + settings: dict = {} + +# Organization endpoints +@app.get("/organizations", response_model=List[Organization]) +async def get_organizations(): + async with db_pool.acquire() as conn: + rows = await conn.fetch(""" + SELECT + id::text, name, description, settings, + created_at::text, updated_at::text + FROM organizations + ORDER BY created_at DESC + """) + + organizations = [] + for row in rows: + organizations.append(Organization( + id=row['id'], + name=row['name'], + description=row['description'], + settings=json.loads(row['settings']) if row['settings'] else {}, + created_at=row['created_at'], + updated_at=row['updated_at'] + )) + + return organizations + +@app.post("/organizations", response_model=Organization) +async def create_organization(org_data: OrganizationCreate): + async with db_pool.acquire() as conn: + org_id = str(uuid.uuid4()) + row = await conn.fetchrow(""" + INSERT INTO organizations (id, name, description, settings) + VALUES ($1, $2, $3, $4) + RETURNING + id::text, name, description, settings, + created_at::text, updated_at::text + """, org_id, org_data.name, org_data.description, json.dumps(org_data.settings)) + + organization = Organization( + id=row['id'], + name=row['name'], + description=row['description'], + settings=json.loads(row['settings']) if row['settings'] else {}, + created_at=row['created_at'], + updated_at=row['updated_at'] + ) + + # Broadcast the change + await manager.broadcast({ + "type": "organization_created", + "data": organization.dict() + }) + + return organization + +@app.get("/organizations/{organization_id}", response_model=Organization) +async def get_organization( + organization_id: str, + user: CurrentUser = Depends(require_user), +): + # SECURITY (issue #6 HIGH-2 / BOLA): authorize the specific org id from the + # path; bare ``WHERE id = $1`` is not an authorization boundary. + require_org_access(organization_id, user) + async with db_pool.acquire() as conn: + row = await conn.fetchrow(""" + SELECT + id::text, name, description, settings, + created_at::text, updated_at::text + FROM organizations + WHERE id = $1 + """, organization_id) + + if not row: + raise HTTPException(status_code=404, detail="Organization not found") + + return Organization( + id=row['id'], + name=row['name'], + description=row['description'], + settings=json.loads(row['settings']) if row['settings'] else {}, + created_at=row['created_at'], + updated_at=row['updated_at'] + ) + +# Team endpoints +@app.get("/teams", response_model=List[Team]) +async def get_teams(organization_id: Optional[str] = None): + async with db_pool.acquire() as conn: + if organization_id: + rows = await conn.fetch(""" + SELECT + id::text, organization_id::text, name, description, + team_type, settings, created_at::text, updated_at::text + FROM teams + WHERE organization_id = $1 + ORDER BY created_at DESC + """, organization_id) + else: + rows = await conn.fetch(""" + SELECT + id::text, organization_id::text, name, description, + team_type, settings, created_at::text, updated_at::text + FROM teams + ORDER BY created_at DESC + """) + + teams = [] + for row in rows: + teams.append(Team( + id=row['id'], + organization_id=row['organization_id'], + name=row['name'], + description=row['description'], + team_type=row['team_type'], + settings=json.loads(row['settings']) if row['settings'] else {}, + created_at=row['created_at'], + updated_at=row['updated_at'] + )) + + return teams + +@app.post("/teams", response_model=Team) +async def create_team(team_data: TeamCreate): + async with db_pool.acquire() as conn: + # Verify organization exists + org_exists = await conn.fetchval( + "SELECT EXISTS(SELECT 1 FROM organizations WHERE id = $1)", + team_data.organization_id + ) + if not org_exists: + raise HTTPException(status_code=404, detail="Organization not found") + + team_id = str(uuid.uuid4()) + row = await conn.fetchrow(""" + INSERT INTO teams (id, organization_id, name, description, team_type, settings) + VALUES ($1, $2, $3, $4, $5, $6) + RETURNING + id::text, organization_id::text, name, description, + team_type, settings, created_at::text, updated_at::text + """, team_id, team_data.organization_id, team_data.name, + team_data.description, team_data.team_type, json.dumps(team_data.settings)) + + team = Team( + id=row['id'], + organization_id=row['organization_id'], + name=row['name'], + description=row['description'], + team_type=row['team_type'], + settings=json.loads(row['settings']) if row['settings'] else {}, + created_at=row['created_at'], + updated_at=row['updated_at'] + ) + + # Broadcast the change + await manager.broadcast({ + "type": "team_created", + "data": team.dict() + }) + + return team + +@app.get("/teams/{team_id}", response_model=Team) +async def get_team( + team_id: str, + user: CurrentUser = Depends(require_user), +): + async with db_pool.acquire() as conn: + row = await conn.fetchrow(""" + SELECT + id::text, organization_id::text, name, description, + team_type, settings, created_at::text, updated_at::text + FROM teams + WHERE id = $1 + """, team_id) + + if not row: + raise HTTPException(status_code=404, detail="Team not found") + # SECURITY (issue #6 HIGH-2): authorize via the team's parent org. + require_org_access(row['organization_id'], user) + return Team( + id=row['id'], + organization_id=row['organization_id'], + name=row['name'], + description=row['description'], + team_type=row['team_type'], + settings=json.loads(row['settings']) if row['settings'] else {}, + created_at=row['created_at'], + updated_at=row['updated_at'] + ) + +# WebSocket endpoint for real-time updates +@app.websocket("/ws") +async def websocket_endpoint(websocket: WebSocket): + # SECURITY (issue #6 CRITICAL-2): authenticate BEFORE accept(). The + # app-wide ``dependencies=[Depends(get_current_user)]`` is a no-op on + # WebSocket routes (WS handshakes have no HTTP response channel); every WS + # handler must call authenticate_websocket() first — matching the pattern + # already used by all handlers in services/orchestrator/main.py. + user = await authenticate_websocket(websocket) + if user is None: + return # authenticate_websocket already closed the socket (1008) + await manager.connect(websocket) + try: + while True: + # Keep the connection alive and listen for client pings + data = await websocket.receive_text() + if data == "ping": + await websocket.send_text("pong") + except WebSocketDisconnect: + manager.disconnect(websocket) + +# Proxy all other requests to the original orchestrator +@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"]) +async def proxy_to_orchestrator(path: str, request: Request): + url = f"{ORCHESTRATOR_URL}/{path}" + + # Handle CORS preflight requests + if request.method == "OPTIONS": + return JSONResponse( + content={}, + headers={ + "Access-Control-Allow-Origin": "*", + "Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, PATCH, OPTIONS", + "Access-Control-Allow-Headers": "*", + } + ) + + try: + async with httpx.AsyncClient(timeout=30.0) as client: + # Get request body if present + body = None + if request.method in ["POST", "PUT", "PATCH"]: + body = await request.body() + # Log the request for debugging + if path == "agents/from-template": + print(f"DEBUG: Proxying agents/from-template request") + print(f"DEBUG: URL: {url}") + print(f"DEBUG: Body: {body.decode() if body else 'None'}") + print(f"DEBUG: Headers: {dict(request.headers)}") + + # Forward the request + response = await client.request( + method=request.method, + url=url, + params=request.query_params, + content=body, + headers={k: v for k, v in request.headers.items() + if k.lower() not in ['host', 'content-length']}, + ) + + # Check if this was an agent creation and broadcast the change + if (request.method == "POST" and + (path == "agents" or path == "agents/from-template") and + response.status_code in [200, 201] and + response.headers.get("content-type", "").startswith("application/json")): + try: + agent_data = response.json() + await manager.broadcast({ + "type": "agent_created", + "data": agent_data + }) + except: + pass # Ignore broadcast errors + + return JSONResponse( + content=response.json() if response.headers.get("content-type", "").startswith("application/json") else response.text, + status_code=response.status_code, + headers={ + "Access-Control-Allow-Origin": "*", + **{k: v for k, v in response.headers.items() + if k.lower() not in ['content-length', 'transfer-encoding', 'connection']} + } + ) + + except httpx.RequestError as e: + print(f"ERROR: RequestError in proxy: {str(e)}") + raise HTTPException(status_code=503, detail=f"Failed to connect to orchestrator: {str(e)}") + except Exception as e: + print(f"ERROR: Exception in proxy: {str(e)}") + import traceback + traceback.print_exc() + raise HTTPException(status_code=500, detail=f"Proxy error: {str(e)}") + +if __name__ == "__main__": + import uvicorn uvicorn.run(app, host="0.0.0.0", port=8006) \ No newline at end of file diff --git a/services/orchestrator/agent_expertise_tracker.py b/services/orchestrator/agent_expertise_tracker.py index 1848811..33009c2 100644 --- a/services/orchestrator/agent_expertise_tracker.py +++ b/services/orchestrator/agent_expertise_tracker.py @@ -1,520 +1,520 @@ -""" -Agent Expertise Tracker - -Provides analytics and insights into agent performance, learning patterns, -and expertise development across the FuzeAgent system. -""" - -import asyncio -import json -import logging -from dataclasses import dataclass -from datetime import datetime, timedelta -from typing import Any, Dict, List, Optional - -from .database import get_db_connection - -logger = logging.getLogger(__name__) - - -@dataclass -class ExpertiseInsight: - """Insight about agent expertise development""" - - agent_id: str - skill_area: str - insight_type: str # 'improving', 'declining', 'plateau', 'breakthrough' - description: str - confidence: float - evidence: Dict[str, Any] - timestamp: datetime - - -@dataclass -class AgentPerformanceMetrics: - """Performance metrics for an agent""" - - agent_id: str - total_tasks: int - success_rate: float - avg_expertise_level: float - improving_skills_count: int - declining_skills_count: int - memory_usage_stats: Dict[str, Any] - recent_performance_trend: str - top_skill_areas: List[Dict[str, Any]] - - -class AgentExpertiseTracker: - """ - Tracks and analyzes agent expertise development, providing insights - into learning patterns, performance trends, and optimization opportunities. - """ - - def __init__(self, database_url: str): - self.database_url = database_url - self.insights_cache: Dict[str, List[ExpertiseInsight]] = {} - self.metrics_cache: Dict[str, AgentPerformanceMetrics] = {} - self.cache_ttl = 300 # 5 minutes - self.last_cache_update = {} - - async def get_agent_performance_metrics( - self, agent_id: str - ) -> Optional[AgentPerformanceMetrics]: - """Get comprehensive performance metrics for an agent""" - - # Check cache first - if ( - agent_id in self.metrics_cache - and agent_id in self.last_cache_update - and (datetime.now() - self.last_cache_update[agent_id]).total_seconds() - < self.cache_ttl - ): - return self.metrics_cache[agent_id] - - try: - async with get_db_connection() as conn: - # Get basic performance stats - basic_stats = await conn.fetchrow( - """ - SELECT - COUNT(DISTINCT am.task_id) as total_tasks, - AVG(CASE WHEN am.memory_type = 'success' THEN 1.0 ELSE 0.0 END) as success_rate, - COUNT(DISTINCT am.id) as total_memories - FROM agent_memory am - WHERE am.agent_id = $1 - """, - agent_id, - ) - - # Get expertise summary - expertise_stats = await conn.fetchrow( - """ - SELECT - AVG(expertise_level) as avg_expertise_level, - COUNT(CASE WHEN performance_trend = 'improving' THEN 1 END) as improving_skills, - COUNT(CASE WHEN performance_trend = 'declining' THEN 1 END) as declining_skills - FROM agent_expertise - WHERE agent_id = $1 - """, - agent_id, - ) - - # Get memory usage statistics - memory_stats = await conn.fetchrow( - """ - SELECT - COUNT(*) as total_memories, - AVG(confidence_score) as avg_confidence, - SUM(usage_count) as total_usage, - COUNT(DISTINCT memory_type) as memory_types_used - FROM agent_memory - WHERE agent_id = $1 - """, - agent_id, - ) - - # Get top skill areas - top_skills = await conn.fetch( - """ - SELECT skill_area, expertise_level, success_rate, task_count, performance_trend - FROM agent_expertise - WHERE agent_id = $1 - ORDER BY expertise_level DESC, success_rate DESC - LIMIT 5 - """, - agent_id, - ) - - # Determine recent performance trend - recent_trend = await self._calculate_recent_trend(agent_id, conn) - - # Build metrics object - metrics = AgentPerformanceMetrics( - agent_id=agent_id, - total_tasks=basic_stats["total_tasks"] or 0, - success_rate=basic_stats["success_rate"] or 0.0, - avg_expertise_level=expertise_stats["avg_expertise_level"] or 0.0, - improving_skills_count=expertise_stats["improving_skills"] or 0, - declining_skills_count=expertise_stats["declining_skills"] or 0, - memory_usage_stats={ - "total_memories": memory_stats["total_memories"] or 0, - "avg_confidence": ( - float(memory_stats["avg_confidence"]) - if memory_stats["avg_confidence"] - else 0.0 - ), - "total_usage": memory_stats["total_usage"] or 0, - "memory_types_used": memory_stats["memory_types_used"] or 0, - }, - recent_performance_trend=recent_trend, - top_skill_areas=[dict(skill) for skill in top_skills], - ) - - # Cache the result - self.metrics_cache[agent_id] = metrics - self.last_cache_update[agent_id] = datetime.now() - - return metrics - - except Exception as e: - logger.error(f"Error getting performance metrics for agent {agent_id}: {e}") - return None - - async def generate_expertise_insights( - self, agent_id: str - ) -> List[ExpertiseInsight]: - """Generate insights about agent expertise development""" - - # Check cache first - if ( - agent_id in self.insights_cache - and agent_id in self.last_cache_update - and (datetime.now() - self.last_cache_update[agent_id]).total_seconds() - < self.cache_ttl - ): - return self.insights_cache[agent_id] - - insights = [] - - try: - async with get_db_connection() as conn: - # Analyze learning velocity patterns - learning_insights = await self._analyze_learning_velocity( - agent_id, conn - ) - insights.extend(learning_insights) - - # Analyze skill development patterns - skill_insights = await self._analyze_skill_development(agent_id, conn) - insights.extend(skill_insights) - - # Analyze memory usage patterns - memory_insights = await self._analyze_memory_patterns(agent_id, conn) - insights.extend(memory_insights) - - # Cache the results - self.insights_cache[agent_id] = insights - self.last_cache_update[agent_id] = datetime.now() - - except Exception as e: - logger.error(f"Error generating insights for agent {agent_id}: {e}") - - return insights - - async def get_system_wide_expertise_summary(self) -> Dict[str, Any]: - """Get system-wide expertise and performance summary""" - - try: - async with get_db_connection() as conn: - # Overall system stats - system_stats = await conn.fetchrow(""" - SELECT - COUNT(DISTINCT a.id) as total_agents, - COUNT(DISTINCT ae.skill_area) as total_skill_areas, - AVG(ae.expertise_level) as avg_system_expertise, - COUNT(CASE WHEN ae.performance_trend = 'improving' THEN 1 END) as improving_agents, - COUNT(CASE WHEN ae.performance_trend = 'declining' THEN 1 END) as declining_agents - FROM agents a - LEFT JOIN agent_expertise ae ON a.id = ae.agent_id - """) - - # Memory system stats - memory_stats = await conn.fetchrow(""" - SELECT - COUNT(*) as total_memories, - AVG(confidence_score) as avg_confidence, - SUM(usage_count) as total_usage, - COUNT(DISTINCT agent_id) as agents_with_memory - FROM agent_memory - """) - - # Top performing skill areas - top_skill_areas = await conn.fetch(""" - SELECT - skill_area, - COUNT(*) as agent_count, - AVG(expertise_level) as avg_expertise, - AVG(success_rate) as avg_success_rate - FROM agent_expertise - GROUP BY skill_area - ORDER BY avg_expertise DESC, avg_success_rate DESC - LIMIT 10 - """) - - # Recent activity - recent_activity = await conn.fetchrow(""" - SELECT - COUNT(CASE WHEN created_at > NOW() - INTERVAL '24 hours' THEN 1 END) as memories_24h, - COUNT(CASE WHEN created_at > NOW() - INTERVAL '7 days' THEN 1 END) as memories_7d, - COUNT(DISTINCT CASE WHEN created_at > NOW() - INTERVAL '24 hours' THEN agent_id END) as active_agents_24h - FROM agent_memory - """) - - return { - "system_stats": dict(system_stats) if system_stats else {}, - "memory_stats": dict(memory_stats) if memory_stats else {}, - "top_skill_areas": [dict(skill) for skill in top_skill_areas], - "recent_activity": dict(recent_activity) if recent_activity else {}, - "timestamp": datetime.now().isoformat(), - } - - except Exception as e: - logger.error(f"Error getting system-wide expertise summary: {e}") - return {"error": str(e)} - - async def _calculate_recent_trend(self, agent_id: str, conn) -> str: - """Calculate recent performance trend for an agent""" - - try: - # Get recent task outcomes - recent_outcomes = await conn.fetch( - """ - SELECT - DATE_TRUNC('day', created_at) as date, - AVG(CASE WHEN memory_type = 'success' THEN 1.0 ELSE 0.0 END) as daily_success_rate - FROM agent_memory - WHERE agent_id = $1 - AND created_at > NOW() - INTERVAL '14 days' - AND memory_type IN ('success', 'task_outcome') - GROUP BY DATE_TRUNC('day', created_at) - ORDER BY date DESC - LIMIT 7 - """, - agent_id, - ) - - if len(recent_outcomes) < 3: - return "insufficient_data" - - # Calculate trend - success_rates = [ - float(row["daily_success_rate"]) for row in recent_outcomes - ] - - # Simple linear trend calculation - if len(success_rates) >= 3: - early_avg = sum(success_rates[-3:]) / 3 - recent_avg = sum(success_rates[:3]) / 3 - - if recent_avg > early_avg + 0.1: - return "improving" - elif recent_avg < early_avg - 0.1: - return "declining" - else: - return "stable" - - return "stable" - - except Exception as e: - logger.error(f"Error calculating recent trend: {e}") - return "unknown" - - async def _analyze_learning_velocity( - self, agent_id: str, conn - ) -> List[ExpertiseInsight]: - """Analyze learning velocity patterns""" - - insights = [] - - try: - # Get skills with high learning velocity - fast_learners = await conn.fetch( - """ - SELECT skill_area, learning_velocity, expertise_level, task_count - FROM agent_expertise - WHERE agent_id = $1 AND learning_velocity > 0.1 - ORDER BY learning_velocity DESC - """, - agent_id, - ) - - for skill in fast_learners: - insights.append( - ExpertiseInsight( - agent_id=agent_id, - skill_area=skill["skill_area"], - insight_type="improving", - description=f"Rapid improvement in {skill['skill_area']} with velocity {skill['learning_velocity']:.2f}", - confidence=0.8, - evidence={ - "learning_velocity": float(skill["learning_velocity"]), - "expertise_level": float(skill["expertise_level"]), - "task_count": skill["task_count"], - }, - timestamp=datetime.now(), - ) - ) - - # Get skills with declining performance - declining_skills = await conn.fetch( - """ - SELECT skill_area, learning_velocity, expertise_level, task_count - FROM agent_expertise - WHERE agent_id = $1 AND learning_velocity < -0.05 - ORDER BY learning_velocity ASC - """, - agent_id, - ) - - for skill in declining_skills: - insights.append( - ExpertiseInsight( - agent_id=agent_id, - skill_area=skill["skill_area"], - insight_type="declining", - description=f"Performance decline in {skill['skill_area']} - may need attention", - confidence=0.7, - evidence={ - "learning_velocity": float(skill["learning_velocity"]), - "expertise_level": float(skill["expertise_level"]), - "task_count": skill["task_count"], - }, - timestamp=datetime.now(), - ) - ) - - except Exception as e: - logger.error(f"Error analyzing learning velocity: {e}") - - return insights - - async def _analyze_skill_development( - self, agent_id: str, conn - ) -> List[ExpertiseInsight]: - """Analyze skill development patterns""" - - insights = [] - - try: - # Find breakthrough moments (significant expertise jumps) - breakthroughs = await conn.fetch( - """ - SELECT skill_area, expertise_level, success_rate, task_count - FROM agent_expertise - WHERE agent_id = $1 - AND expertise_level > 0.7 - AND success_rate > 0.8 - AND task_count >= 5 - """, - agent_id, - ) - - for breakthrough in breakthroughs: - insights.append( - ExpertiseInsight( - agent_id=agent_id, - skill_area=breakthrough["skill_area"], - insight_type="breakthrough", - description=f"Expert level achieved in {breakthrough['skill_area']} with {breakthrough['success_rate']:.1%} success rate", - confidence=0.9, - evidence={ - "expertise_level": float(breakthrough["expertise_level"]), - "success_rate": float(breakthrough["success_rate"]), - "task_count": breakthrough["task_count"], - }, - timestamp=datetime.now(), - ) - ) - - # Find plateau situations (high task count but low expertise) - plateaus = await conn.fetch( - """ - SELECT skill_area, expertise_level, success_rate, task_count - FROM agent_expertise - WHERE agent_id = $1 - AND task_count > 10 - AND expertise_level < 0.4 - AND learning_velocity BETWEEN -0.02 AND 0.02 - """, - agent_id, - ) - - for plateau in plateaus: - insights.append( - ExpertiseInsight( - agent_id=agent_id, - skill_area=plateau["skill_area"], - insight_type="plateau", - description=f"Learning plateau in {plateau['skill_area']} - consider new approaches", - confidence=0.6, - evidence={ - "expertise_level": float(plateau["expertise_level"]), - "success_rate": float(plateau["success_rate"]), - "task_count": plateau["task_count"], - }, - timestamp=datetime.now(), - ) - ) - - except Exception as e: - logger.error(f"Error analyzing skill development: {e}") - - return insights - - async def _analyze_memory_patterns( - self, agent_id: str, conn - ) -> List[ExpertiseInsight]: - """Analyze memory usage and effectiveness patterns""" - - insights = [] - - try: - # Analyze memory types and their effectiveness - memory_effectiveness = await conn.fetchrow( - """ - SELECT - COUNT(*) as total_memories, - AVG(usage_count) as avg_usage, - AVG(confidence_score) as avg_confidence, - COUNT(CASE WHEN usage_count > 5 THEN 1 END) as high_usage_memories - FROM agent_memory - WHERE agent_id = $1 - """, - agent_id, - ) - - if memory_effectiveness and memory_effectiveness["total_memories"] > 50: - high_usage_ratio = ( - memory_effectiveness["high_usage_memories"] - / memory_effectiveness["total_memories"] - ) - - if ( - high_usage_ratio > 0.2 - ): # More than 20% of memories are highly reused - insights.append( - ExpertiseInsight( - agent_id=agent_id, - skill_area="memory_management", - insight_type="improving", - description=f"Excellent memory reuse patterns - {high_usage_ratio:.1%} of memories are frequently accessed", - confidence=0.8, - evidence={ - "total_memories": memory_effectiveness[ - "total_memories" - ], - "high_usage_ratio": high_usage_ratio, - "avg_confidence": float( - memory_effectiveness["avg_confidence"] - ), - }, - timestamp=datetime.now(), - ) - ) - - except Exception as e: - logger.error(f"Error analyzing memory patterns: {e}") - - return insights - - async def clear_cache(self, agent_id: Optional[str] = None): - """Clear analytics cache""" - if agent_id: - self.insights_cache.pop(agent_id, None) - self.metrics_cache.pop(agent_id, None) - self.last_cache_update.pop(agent_id, None) - else: - self.insights_cache.clear() - self.metrics_cache.clear() - self.last_cache_update.clear() +""" +Agent Expertise Tracker + +Provides analytics and insights into agent performance, learning patterns, +and expertise development across the FuzeAgent system. +""" + +import asyncio +import json +import logging +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Any, Dict, List, Optional + +from .database import get_db_connection + +logger = logging.getLogger(__name__) + + +@dataclass +class ExpertiseInsight: + """Insight about agent expertise development""" + + agent_id: str + skill_area: str + insight_type: str # 'improving', 'declining', 'plateau', 'breakthrough' + description: str + confidence: float + evidence: Dict[str, Any] + timestamp: datetime + + +@dataclass +class AgentPerformanceMetrics: + """Performance metrics for an agent""" + + agent_id: str + total_tasks: int + success_rate: float + avg_expertise_level: float + improving_skills_count: int + declining_skills_count: int + memory_usage_stats: Dict[str, Any] + recent_performance_trend: str + top_skill_areas: List[Dict[str, Any]] + + +class AgentExpertiseTracker: + """ + Tracks and analyzes agent expertise development, providing insights + into learning patterns, performance trends, and optimization opportunities. + """ + + def __init__(self, database_url: str): + self.database_url = database_url + self.insights_cache: Dict[str, List[ExpertiseInsight]] = {} + self.metrics_cache: Dict[str, AgentPerformanceMetrics] = {} + self.cache_ttl = 300 # 5 minutes + self.last_cache_update = {} + + async def get_agent_performance_metrics( + self, agent_id: str + ) -> Optional[AgentPerformanceMetrics]: + """Get comprehensive performance metrics for an agent""" + + # Check cache first + if ( + agent_id in self.metrics_cache + and agent_id in self.last_cache_update + and (datetime.now() - self.last_cache_update[agent_id]).total_seconds() + < self.cache_ttl + ): + return self.metrics_cache[agent_id] + + try: + async with get_db_connection() as conn: + # Get basic performance stats + basic_stats = await conn.fetchrow( + """ + SELECT + COUNT(DISTINCT am.task_id) as total_tasks, + AVG(CASE WHEN am.memory_type = 'success' THEN 1.0 ELSE 0.0 END) as success_rate, + COUNT(DISTINCT am.id) as total_memories + FROM agent_memory am + WHERE am.agent_id = $1 + """, + agent_id, + ) + + # Get expertise summary + expertise_stats = await conn.fetchrow( + """ + SELECT + AVG(expertise_level) as avg_expertise_level, + COUNT(CASE WHEN performance_trend = 'improving' THEN 1 END) as improving_skills, + COUNT(CASE WHEN performance_trend = 'declining' THEN 1 END) as declining_skills + FROM agent_expertise + WHERE agent_id = $1 + """, + agent_id, + ) + + # Get memory usage statistics + memory_stats = await conn.fetchrow( + """ + SELECT + COUNT(*) as total_memories, + AVG(confidence_score) as avg_confidence, + SUM(usage_count) as total_usage, + COUNT(DISTINCT memory_type) as memory_types_used + FROM agent_memory + WHERE agent_id = $1 + """, + agent_id, + ) + + # Get top skill areas + top_skills = await conn.fetch( + """ + SELECT skill_area, expertise_level, success_rate, task_count, performance_trend + FROM agent_expertise + WHERE agent_id = $1 + ORDER BY expertise_level DESC, success_rate DESC + LIMIT 5 + """, + agent_id, + ) + + # Determine recent performance trend + recent_trend = await self._calculate_recent_trend(agent_id, conn) + + # Build metrics object + metrics = AgentPerformanceMetrics( + agent_id=agent_id, + total_tasks=basic_stats["total_tasks"] or 0, + success_rate=basic_stats["success_rate"] or 0.0, + avg_expertise_level=expertise_stats["avg_expertise_level"] or 0.0, + improving_skills_count=expertise_stats["improving_skills"] or 0, + declining_skills_count=expertise_stats["declining_skills"] or 0, + memory_usage_stats={ + "total_memories": memory_stats["total_memories"] or 0, + "avg_confidence": ( + float(memory_stats["avg_confidence"]) + if memory_stats["avg_confidence"] + else 0.0 + ), + "total_usage": memory_stats["total_usage"] or 0, + "memory_types_used": memory_stats["memory_types_used"] or 0, + }, + recent_performance_trend=recent_trend, + top_skill_areas=[dict(skill) for skill in top_skills], + ) + + # Cache the result + self.metrics_cache[agent_id] = metrics + self.last_cache_update[agent_id] = datetime.now() + + return metrics + + except Exception as e: + logger.error(f"Error getting performance metrics for agent {agent_id}: {e}") + return None + + async def generate_expertise_insights( + self, agent_id: str + ) -> List[ExpertiseInsight]: + """Generate insights about agent expertise development""" + + # Check cache first + if ( + agent_id in self.insights_cache + and agent_id in self.last_cache_update + and (datetime.now() - self.last_cache_update[agent_id]).total_seconds() + < self.cache_ttl + ): + return self.insights_cache[agent_id] + + insights = [] + + try: + async with get_db_connection() as conn: + # Analyze learning velocity patterns + learning_insights = await self._analyze_learning_velocity( + agent_id, conn + ) + insights.extend(learning_insights) + + # Analyze skill development patterns + skill_insights = await self._analyze_skill_development(agent_id, conn) + insights.extend(skill_insights) + + # Analyze memory usage patterns + memory_insights = await self._analyze_memory_patterns(agent_id, conn) + insights.extend(memory_insights) + + # Cache the results + self.insights_cache[agent_id] = insights + self.last_cache_update[agent_id] = datetime.now() + + except Exception as e: + logger.error(f"Error generating insights for agent {agent_id}: {e}") + + return insights + + async def get_system_wide_expertise_summary(self) -> Dict[str, Any]: + """Get system-wide expertise and performance summary""" + + try: + async with get_db_connection() as conn: + # Overall system stats + system_stats = await conn.fetchrow(""" + SELECT + COUNT(DISTINCT a.id) as total_agents, + COUNT(DISTINCT ae.skill_area) as total_skill_areas, + AVG(ae.expertise_level) as avg_system_expertise, + COUNT(CASE WHEN ae.performance_trend = 'improving' THEN 1 END) as improving_agents, + COUNT(CASE WHEN ae.performance_trend = 'declining' THEN 1 END) as declining_agents + FROM agents a + LEFT JOIN agent_expertise ae ON a.id = ae.agent_id + """) + + # Memory system stats + memory_stats = await conn.fetchrow(""" + SELECT + COUNT(*) as total_memories, + AVG(confidence_score) as avg_confidence, + SUM(usage_count) as total_usage, + COUNT(DISTINCT agent_id) as agents_with_memory + FROM agent_memory + """) + + # Top performing skill areas + top_skill_areas = await conn.fetch(""" + SELECT + skill_area, + COUNT(*) as agent_count, + AVG(expertise_level) as avg_expertise, + AVG(success_rate) as avg_success_rate + FROM agent_expertise + GROUP BY skill_area + ORDER BY avg_expertise DESC, avg_success_rate DESC + LIMIT 10 + """) + + # Recent activity + recent_activity = await conn.fetchrow(""" + SELECT + COUNT(CASE WHEN created_at > NOW() - INTERVAL '24 hours' THEN 1 END) as memories_24h, + COUNT(CASE WHEN created_at > NOW() - INTERVAL '7 days' THEN 1 END) as memories_7d, + COUNT(DISTINCT CASE WHEN created_at > NOW() - INTERVAL '24 hours' THEN agent_id END) as active_agents_24h + FROM agent_memory + """) + + return { + "system_stats": dict(system_stats) if system_stats else {}, + "memory_stats": dict(memory_stats) if memory_stats else {}, + "top_skill_areas": [dict(skill) for skill in top_skill_areas], + "recent_activity": dict(recent_activity) if recent_activity else {}, + "timestamp": datetime.now().isoformat(), + } + + except Exception as e: + logger.error(f"Error getting system-wide expertise summary: {e}") + return {"error": str(e)} + + async def _calculate_recent_trend(self, agent_id: str, conn) -> str: + """Calculate recent performance trend for an agent""" + + try: + # Get recent task outcomes + recent_outcomes = await conn.fetch( + """ + SELECT + DATE_TRUNC('day', created_at) as date, + AVG(CASE WHEN memory_type = 'success' THEN 1.0 ELSE 0.0 END) as daily_success_rate + FROM agent_memory + WHERE agent_id = $1 + AND created_at > NOW() - INTERVAL '14 days' + AND memory_type IN ('success', 'task_outcome') + GROUP BY DATE_TRUNC('day', created_at) + ORDER BY date DESC + LIMIT 7 + """, + agent_id, + ) + + if len(recent_outcomes) < 3: + return "insufficient_data" + + # Calculate trend + success_rates = [ + float(row["daily_success_rate"]) for row in recent_outcomes + ] + + # Simple linear trend calculation + if len(success_rates) >= 3: + early_avg = sum(success_rates[-3:]) / 3 + recent_avg = sum(success_rates[:3]) / 3 + + if recent_avg > early_avg + 0.1: + return "improving" + elif recent_avg < early_avg - 0.1: + return "declining" + else: + return "stable" + + return "stable" + + except Exception as e: + logger.error(f"Error calculating recent trend: {e}") + return "unknown" + + async def _analyze_learning_velocity( + self, agent_id: str, conn + ) -> List[ExpertiseInsight]: + """Analyze learning velocity patterns""" + + insights = [] + + try: + # Get skills with high learning velocity + fast_learners = await conn.fetch( + """ + SELECT skill_area, learning_velocity, expertise_level, task_count + FROM agent_expertise + WHERE agent_id = $1 AND learning_velocity > 0.1 + ORDER BY learning_velocity DESC + """, + agent_id, + ) + + for skill in fast_learners: + insights.append( + ExpertiseInsight( + agent_id=agent_id, + skill_area=skill["skill_area"], + insight_type="improving", + description=f"Rapid improvement in {skill['skill_area']} with velocity {skill['learning_velocity']:.2f}", + confidence=0.8, + evidence={ + "learning_velocity": float(skill["learning_velocity"]), + "expertise_level": float(skill["expertise_level"]), + "task_count": skill["task_count"], + }, + timestamp=datetime.now(), + ) + ) + + # Get skills with declining performance + declining_skills = await conn.fetch( + """ + SELECT skill_area, learning_velocity, expertise_level, task_count + FROM agent_expertise + WHERE agent_id = $1 AND learning_velocity < -0.05 + ORDER BY learning_velocity ASC + """, + agent_id, + ) + + for skill in declining_skills: + insights.append( + ExpertiseInsight( + agent_id=agent_id, + skill_area=skill["skill_area"], + insight_type="declining", + description=f"Performance decline in {skill['skill_area']} - may need attention", + confidence=0.7, + evidence={ + "learning_velocity": float(skill["learning_velocity"]), + "expertise_level": float(skill["expertise_level"]), + "task_count": skill["task_count"], + }, + timestamp=datetime.now(), + ) + ) + + except Exception as e: + logger.error(f"Error analyzing learning velocity: {e}") + + return insights + + async def _analyze_skill_development( + self, agent_id: str, conn + ) -> List[ExpertiseInsight]: + """Analyze skill development patterns""" + + insights = [] + + try: + # Find breakthrough moments (significant expertise jumps) + breakthroughs = await conn.fetch( + """ + SELECT skill_area, expertise_level, success_rate, task_count + FROM agent_expertise + WHERE agent_id = $1 + AND expertise_level > 0.7 + AND success_rate > 0.8 + AND task_count >= 5 + """, + agent_id, + ) + + for breakthrough in breakthroughs: + insights.append( + ExpertiseInsight( + agent_id=agent_id, + skill_area=breakthrough["skill_area"], + insight_type="breakthrough", + description=f"Expert level achieved in {breakthrough['skill_area']} with {breakthrough['success_rate']:.1%} success rate", + confidence=0.9, + evidence={ + "expertise_level": float(breakthrough["expertise_level"]), + "success_rate": float(breakthrough["success_rate"]), + "task_count": breakthrough["task_count"], + }, + timestamp=datetime.now(), + ) + ) + + # Find plateau situations (high task count but low expertise) + plateaus = await conn.fetch( + """ + SELECT skill_area, expertise_level, success_rate, task_count + FROM agent_expertise + WHERE agent_id = $1 + AND task_count > 10 + AND expertise_level < 0.4 + AND learning_velocity BETWEEN -0.02 AND 0.02 + """, + agent_id, + ) + + for plateau in plateaus: + insights.append( + ExpertiseInsight( + agent_id=agent_id, + skill_area=plateau["skill_area"], + insight_type="plateau", + description=f"Learning plateau in {plateau['skill_area']} - consider new approaches", + confidence=0.6, + evidence={ + "expertise_level": float(plateau["expertise_level"]), + "success_rate": float(plateau["success_rate"]), + "task_count": plateau["task_count"], + }, + timestamp=datetime.now(), + ) + ) + + except Exception as e: + logger.error(f"Error analyzing skill development: {e}") + + return insights + + async def _analyze_memory_patterns( + self, agent_id: str, conn + ) -> List[ExpertiseInsight]: + """Analyze memory usage and effectiveness patterns""" + + insights = [] + + try: + # Analyze memory types and their effectiveness + memory_effectiveness = await conn.fetchrow( + """ + SELECT + COUNT(*) as total_memories, + AVG(usage_count) as avg_usage, + AVG(confidence_score) as avg_confidence, + COUNT(CASE WHEN usage_count > 5 THEN 1 END) as high_usage_memories + FROM agent_memory + WHERE agent_id = $1 + """, + agent_id, + ) + + if memory_effectiveness and memory_effectiveness["total_memories"] > 50: + high_usage_ratio = ( + memory_effectiveness["high_usage_memories"] + / memory_effectiveness["total_memories"] + ) + + if ( + high_usage_ratio > 0.2 + ): # More than 20% of memories are highly reused + insights.append( + ExpertiseInsight( + agent_id=agent_id, + skill_area="memory_management", + insight_type="improving", + description=f"Excellent memory reuse patterns - {high_usage_ratio:.1%} of memories are frequently accessed", + confidence=0.8, + evidence={ + "total_memories": memory_effectiveness[ + "total_memories" + ], + "high_usage_ratio": high_usage_ratio, + "avg_confidence": float( + memory_effectiveness["avg_confidence"] + ), + }, + timestamp=datetime.now(), + ) + ) + + except Exception as e: + logger.error(f"Error analyzing memory patterns: {e}") + + return insights + + async def clear_cache(self, agent_id: Optional[str] = None): + """Clear analytics cache""" + if agent_id: + self.insights_cache.pop(agent_id, None) + self.metrics_cache.pop(agent_id, None) + self.last_cache_update.pop(agent_id, None) + else: + self.insights_cache.clear() + self.metrics_cache.clear() + self.last_cache_update.clear() diff --git a/services/orchestrator/claude_code_wrapper.py b/services/orchestrator/claude_code_wrapper.py index 5188ab8..f1186e5 100644 --- a/services/orchestrator/claude_code_wrapper.py +++ b/services/orchestrator/claude_code_wrapper.py @@ -1,842 +1,842 @@ -import asyncio -import json -import os -import subprocess # nosec B404 -- used only with static executable names + arg lists and shell=False (see _run_tests) -import tempfile -import time -from pathlib import Path -from typing import Any, Dict, List, Optional, Type - -# Import Anthropic SDK for real Claude integration -import anthropic -from anthropic import Anthropic -from crewai.tools import BaseTool -from pydantic import BaseModel, Field - -# Import conversation manager for full chat tracking. -# This module is imported both as part of the `services.orchestrator` package -# (relative form, e.g. from main.py/agent_manager.py) and flat with -# services/orchestrator on sys.path (e.g. from tests and main_with_hierarchy.py), -# so support both — mirrors the existing pattern in hierarchy_endpoints.py. -try: - from .conversation_manager import ConversationManager, MessageType -except ImportError: # pragma: no cover - flat import (no parent package) - from conversation_manager import ConversationManager, MessageType - - -class ClaudeCodeInput(BaseModel): - """Input schema for Claude Code tool""" - - task: str = Field(description="Coding task to complete") - language: str = Field(default="python", description="Programming language") - context: str = Field(default="", description="Additional context or requirements") - include_tests: bool = Field( - default=True, description="Whether to include unit tests" - ) - include_docs: bool = Field( - default=True, description="Whether to include documentation" - ) - file_path: Optional[str] = Field( - default=None, description="Optional file path for code context" - ) - - -class ClaudeCodeWrapper(BaseTool): - name: str = "claude_code" - description: str = """ - Execute advanced coding tasks using Claude AI with real-time code generation, - testing, and documentation. Supports multiple programming languages and - follows industry best practices. Enhanced for repository context and Git integration. - """ - args_schema: Type[BaseModel] = ClaudeCodeInput - - # Runtime attributes. ``BaseTool`` is a Pydantic v2 model, which rejects - # assignment to undeclared attributes ("object has no field ..."). These - # are declared as model fields (rather than PrivateAttr) so they remain - # publicly readable on the instance (e.g. ``wrapper.client`` / - # ``wrapper.model``), preserving the tool's public interface. Object - # handles (Anthropic SDK client, git/conversation managers) are typed - # ``Any`` so Pydantic stores them as-is without schema validation. - client: Optional[Any] = None - model: str = "claude-3-5-sonnet-20241022" - workspace_path: Optional[str] = None - git_manager: Optional[Any] = None - agent_id: Optional[str] = None - task_id: Optional[str] = None - conversation_manager: Optional[Any] = None - conversation_session_id: Optional[str] = None - current_context: Dict[str, Any] = Field(default_factory=dict) - repository_context: Dict[str, Any] = Field(default_factory=dict) - - def __init__( - self, - workspace_path: Optional[str] = None, - git_manager: Optional[Any] = None, - agent_id: Optional[str] = None, - task_id: Optional[str] = None, - conversation_manager: Optional[ConversationManager] = None, - ): - super().__init__() - self.client = Anthropic( - api_key=os.getenv("ANTHROPIC_API_KEY"), - ) - self.model = "claude-3-5-sonnet-20241022" - self.workspace_path = workspace_path or os.getcwd() - self.git_manager = git_manager - self.agent_id = agent_id - self.task_id = task_id - self.conversation_manager = conversation_manager or ConversationManager() - self.conversation_session_id: Optional[str] = None - self.current_context = {} # Store context between iterations - - # Repository context - self.repository_context = { - "files_changed": [], - "current_branch": None, - "last_commit": None, - "iteration_count": 0, - } - - def _run( - self, - task: str, - language: str = "python", - context: str = "", - include_tests: bool = True, - include_docs: bool = True, - file_path: Optional[str] = None, - iteration_number: Optional[int] = None, - ) -> str: - """Execute Claude Code for a specific task with real AI integration""" - - try: - # Update iteration count - if iteration_number: - self.repository_context["iteration_count"] = iteration_number - else: - self.repository_context["iteration_count"] += 1 - - # Get repository context if Git manager is available - repo_context = "" - if self.git_manager: - try: - # Note: In a full implementation, we'd make this method async - # For now, we'll skip the Git context in the sync version - repo_context = ( - "Repository context: Available (Git manager configured)" - ) - except Exception as e: - repo_context = f"Repository context unavailable: {str(e)}" - - # Read existing file context if provided - existing_code = "" - if file_path: - # Use workspace-relative path if available - full_path = ( - os.path.join(self.workspace_path, file_path) - if not os.path.isabs(file_path) - else file_path - ) - if os.path.exists(full_path): - with open(full_path, "r") as f: - existing_code = f.read() - - # Prepare the comprehensive prompt with repository context - prompt = self._build_prompt( - task=task, - language=language, - context=context, - existing_code=existing_code, - include_tests=include_tests, - include_docs=include_docs, - repo_context=repo_context, - ) - - # Call Claude API - response = self.client.messages.create( - model=self.model, - max_tokens=4096, - temperature=0.3, # Lower temperature for more consistent code - messages=[{"role": "user", "content": prompt}], - ) - - # Parse the response and extract code files - result = self._parse_response( - response.content[0].text, language, include_tests, include_docs - ) - - # Save files to workspace if available, otherwise use temp directory - if self.workspace_path and os.path.exists(self.workspace_path): - saved_files = self._save_files_to_workspace(result["files"]) - else: - with tempfile.TemporaryDirectory() as tmpdir: - saved_files = self._save_files(result["files"], tmpdir) - - # Update repository context with changed files - self.repository_context["files_changed"].extend( - [ - f["filename"] - for f in result["files"] - if f["type"] == "implementation" - ] - ) - - # Run tests if generated and in workspace - test_results = None - if include_tests and any(f["type"] == "test" for f in result["files"]): - if self.workspace_path and os.path.exists(self.workspace_path): - test_results = self._run_tests(self.workspace_path, language) - else: - with tempfile.TemporaryDirectory() as tmpdir: - self._save_files(result["files"], tmpdir) - test_results = self._run_tests(tmpdir, language) - - return json.dumps( - { - "status": "success", - "files": result["files"], - "explanation": result.get("explanation", ""), - "test_results": test_results, - "commit_message": result.get("commit_message", ""), - "execution_summary": f"Generated {len(result['files'])} files for {language} task: {task[:100]}...", - "iteration": self.repository_context["iteration_count"], - "workspace_path": self.workspace_path, - "repository_context": self.repository_context, - } - ) - - except anthropic.APIError as e: - return json.dumps( - { - "status": "error", - "error": f"Claude API error: {str(e)}", - "error_type": "api_error", - } - ) - except Exception as e: - return json.dumps( - { - "status": "error", - "error": f"Unexpected error: {str(e)}", - "error_type": "general_error", - } - ) - - def _build_prompt( - self, - task: str, - language: str, - context: str, - existing_code: str, - include_tests: bool, - include_docs: bool, - repo_context: str = "", - ) -> str: - """Build a comprehensive prompt for Claude with repository context""" - - # Build agent context - agent_info = "" - if self.agent_id and self.task_id: - agent_info = f""" -**Agent Context**: -- Agent ID: {self.agent_id} -- Task ID: {self.task_id} -- Iteration: {self.repository_context['iteration_count']} -- Workspace: {self.workspace_path} -""" - - prompt = f""" -You are an expert {language} developer working autonomously as part of FuzeAgent AI team. I need you to complete the following coding task: - -{agent_info} - -**Task**: {task} - -**Programming Language**: {language} - -**Additional Context**: {context} - -{repo_context} - -**Existing Code** (if any): -```{language} -{existing_code} -``` - -**Requirements**: -1. Write clean, maintainable, and well-documented code -2. Follow {language} best practices and conventions -3. Include proper error handling -4. Use type hints (where applicable) -5. {"Include comprehensive unit tests" if include_tests else "Focus only on implementation"} -6. {"Include docstrings and comments" if include_docs else "Minimal documentation"} -7. Consider the repository context and maintain consistency with existing code -8. Write code that integrates well with the current branch and recent changes - -**Output Format**: -Please structure your response as follows: - -## Explanation -Brief explanation of your approach and key decisions, considering the repository context. - -## Implementation - -### Main Code -```{language} -# Your main implementation here -``` - -{"### Tests" if include_tests else ""} -{f"```{language}" if include_tests else ""} -{"# Your test code here" if include_tests else ""} -{f"```" if include_tests else ""} - -{"### Documentation" if include_docs else ""} -{"```markdown" if include_docs else ""} -{"# Your documentation here" if include_docs else ""} -{f"```" if include_docs else ""} - -## Commit Message -Suggest a concise git commit message for these changes that follows the repository's commit history style. - -Please ensure the code is production-ready and follows industry standards. -""" - return prompt - - def _parse_response( - self, response: str, language: str, include_tests: bool, include_docs: bool - ) -> Dict[str, Any]: - """Parse Claude's response and extract code files""" - - files = [] - explanation = "" - commit_message = "" - - # Extract explanation - if "## Explanation" in response: - explanation_start = response.find("## Explanation") + len("## Explanation") - explanation_end = response.find("## Implementation") - if explanation_end > explanation_start: - explanation = response[explanation_start:explanation_end].strip() - - # Extract commit message - if "## Commit Message" in response: - commit_start = response.find("## Commit Message") + len("## Commit Message") - commit_message = response[commit_start:].strip() - # Clean up the commit message - commit_message = commit_message.split("\n")[0].strip() - - # Extract main code - main_code = self._extract_code_block(response, "### Main Code", language) - if main_code: - file_ext = self._get_file_extension(language) - files.append( - { - "filename": f"main.{file_ext}", - "content": main_code, - "type": "implementation", - "language": language, - } - ) - - # Extract tests if requested - if include_tests: - test_code = self._extract_code_block(response, "### Tests", language) - if test_code: - test_ext = self._get_file_extension(language) - files.append( - { - "filename": f"test_main.{test_ext}", - "content": test_code, - "type": "test", - "language": language, - } - ) - - # Extract documentation if requested - if include_docs: - docs = self._extract_code_block(response, "### Documentation", "markdown") - if docs: - files.append( - { - "filename": "README.md", - "content": docs, - "type": "documentation", - "language": "markdown", - } - ) - - return { - "files": files, - "explanation": explanation, - "commit_message": commit_message, - } - - def _extract_code_block( - self, text: str, section: str, language: str - ) -> Optional[str]: - """Extract code block from a specific section""" - - section_start = text.find(section) - if section_start == -1: - return None - - # Find the start of the code block - code_start = text.find(f"```{language}", section_start) - if code_start == -1: - code_start = text.find("```", section_start) - if code_start == -1: - return None - - # Find the end of the code block - code_content_start = text.find("\n", code_start) + 1 - code_end = text.find("```", code_content_start) - - if code_end == -1: - return None - - return text[code_content_start:code_end].strip() - - def _get_file_extension(self, language: str) -> str: - """Get appropriate file extension for language""" - extensions = { - "python": "py", - "javascript": "js", - "typescript": "ts", - "java": "java", - "cpp": "cpp", - "c": "c", - "rust": "rs", - "go": "go", - "ruby": "rb", - "php": "php", - "swift": "swift", - "kotlin": "kt", - "scala": "scala", - "r": "R", - "sql": "sql", - "html": "html", - "css": "css", - "shell": "sh", - "bash": "sh", - } - return extensions.get(language.lower(), "txt") - - def _save_files(self, files: List[Dict], tmpdir: str) -> List[str]: - """Save generated files to temporary directory""" - saved_files = [] - - for file_info in files: - file_path = os.path.join(tmpdir, file_info["filename"]) - with open(file_path, "w") as f: - f.write(file_info["content"]) - saved_files.append(file_path) - - return saved_files - - def _run_tests(self, tmpdir: str, language: str) -> Optional[Dict[str, Any]]: - """Run tests for the generated code""" - - try: - if language == "python": - # Try to run pytest - result = subprocess.run( # nosec B603 B607 -- shell=False, static argv; runs pytest on generated code inside an isolated tmpdir - ["python", "-m", "pytest", tmpdir, "-v"], - capture_output=True, - text=True, - timeout=60, - cwd=tmpdir, - ) - - return { - "exit_code": result.returncode, - "stdout": result.stdout, - "stderr": result.stderr, - "success": result.returncode == 0, - } - elif language == "javascript": - # Try to run with node - test_files = [f for f in os.listdir(tmpdir) if f.startswith("test_")] - if test_files: - result = subprocess.run( # nosec B603 B607 -- shell=False, static argv; runs a generated test file inside an isolated tmpdir - ["node", test_files[0]], - capture_output=True, - text=True, - timeout=60, - cwd=tmpdir, - ) - - return { - "exit_code": result.returncode, - "stdout": result.stdout, - "stderr": result.stderr, - "success": result.returncode == 0, - } - - return None - - except subprocess.TimeoutExpired: - return { - "exit_code": -1, - "stdout": "", - "stderr": "Test execution timed out", - "success": False, - } - except Exception as e: - return { - "exit_code": -1, - "stdout": "", - "stderr": f"Test execution error: {str(e)}", - "success": False, - } - - def _build_repository_context( - self, branch_status: Dict[str, Any], commit_history: List[Any] - ) -> str: - """Build repository context string for the prompt""" - context_parts = ["**Repository Context**:"] - - if branch_status: - current_branch = branch_status.get("current_branch", "unknown") - feature_branch = branch_status.get("feature_branch") - has_changes = branch_status.get("has_uncommitted_changes", False) - remote_status = branch_status.get("remote_status", "unknown") - - context_parts.append(f"- Current Branch: `{current_branch}`") - if feature_branch: - context_parts.append(f"- Feature Branch: `{feature_branch}`") - context_parts.append( - f"- Uncommitted Changes: {'Yes' if has_changes else 'No'}" - ) - context_parts.append(f"- Remote Status: {remote_status}") - - if commit_history: - context_parts.append("- Recent Commits:") - for i, commit in enumerate(commit_history[:3]): - context_parts.append(f" {i+1}. `{commit.hash[:8]}` - {commit.message}") - if commit.files_changed: - context_parts.append( - f" Files: {', '.join(commit.files_changed[:5])}" - ) - - if self.repository_context.get("files_changed"): - changed_files = list(set(self.repository_context["files_changed"])) - context_parts.append( - f"- Files Modified This Session: {', '.join(changed_files)}" - ) - - return "\n".join(context_parts) + "\n" - - def _save_files_to_workspace(self, files: List[Dict]) -> List[str]: - """Save generated files directly to workspace""" - saved_files = [] - - for file_info in files: - file_path = os.path.join(self.workspace_path, file_info["filename"]) - - # Create directory if needed - os.makedirs(os.path.dirname(file_path), exist_ok=True) - - with open(file_path, "w") as f: - f.write(file_info["content"]) - saved_files.append(file_path) - - return saved_files - - async def commit_and_push_changes( - self, commit_message: str, files: Optional[List[str]] = None - ) -> Dict[str, Any]: - """Commit and push changes using Git manager""" - if not self.git_manager: - return {"success": False, "error": "No Git manager available"} - - try: - # Commit changes - commit_hash = await self.git_manager.commit_changes( - message=commit_message, - files=files, - iteration_number=self.repository_context["iteration_count"], - ) - - if commit_hash: - self.repository_context["last_commit"] = commit_message - return { - "success": True, - "commit_hash": commit_hash, - "message": "Changes committed successfully", - } - else: - return {"success": True, "message": "No changes to commit"} - - except Exception as e: - return {"success": False, "error": f"Failed to commit changes: {str(e)}"} - - def get_repository_context(self) -> Dict[str, Any]: - """Get current repository context""" - return self.repository_context.copy() - - def reset_context(self): - """Reset the repository context""" - self.repository_context = { - "files_changed": [], - "current_branch": None, - "last_commit": None, - "iteration_count": 0, - } - - async def start_conversation_session(self, sandbox_id: str) -> str: - """Start a conversation session for tracking all Claude Code interactions""" - if not self.agent_id or not self.task_id: - raise ValueError("Agent ID and Task ID required for conversation tracking") - - self.conversation_session_id = ( - await self.conversation_manager.start_conversation_session( - agent_id=self.agent_id, - task_id=self.task_id, - sandbox_id=sandbox_id, - metadata={ - "workspace_path": self.workspace_path, - "model": self.model, - "git_enabled": bool(self.git_manager), - }, - ) - ) - return self.conversation_session_id - - async def end_conversation_session(self) -> bool: - """End the current conversation session""" - if not self.conversation_session_id: - return False - - success = await self.conversation_manager.end_conversation_session( - self.conversation_session_id - ) - self.conversation_session_id = None - return success - - async def execute_task_async( - self, - task: str, - language: str = "python", - context: str = "", - include_tests: bool = True, - include_docs: bool = True, - file_path: Optional[str] = None, - iteration_number: Optional[int] = None, - ) -> Dict[str, Any]: - """Async version of task execution with full conversation tracking""" - - try: - # Update iteration count - if iteration_number: - self.repository_context["iteration_count"] = iteration_number - else: - self.repository_context["iteration_count"] += 1 - - current_iteration = self.repository_context["iteration_count"] - - # Get repository context if Git manager is available - repo_context = "" - if self.git_manager: - try: - branch_status = await self.git_manager.get_branch_status() - self.repository_context["current_branch"] = branch_status.get( - "current_branch" - ) - - commit_history = await self.git_manager.get_commit_history(limit=3) - if commit_history: - self.repository_context["last_commit"] = commit_history[ - 0 - ].message - - repo_context = self._build_repository_context( - branch_status, commit_history - ) - except Exception as e: - repo_context = f"Repository context unavailable: {str(e)}" - - # Read existing file context if provided - existing_code = "" - if file_path: - # Use workspace-relative path if available - full_path = ( - os.path.join(self.workspace_path, file_path) - if not os.path.isabs(file_path) - else file_path - ) - if os.path.exists(full_path): - with open(full_path, "r") as f: - existing_code = f.read() - - # Prepare the comprehensive prompt with repository context - prompt = self._build_prompt( - task=task, - language=language, - context=context, - existing_code=existing_code, - include_tests=include_tests, - include_docs=include_docs, - repo_context=repo_context, - ) - - # Store user prompt in conversation history - if self.conversation_session_id and self.task_id: - await self.conversation_manager.store_user_prompt( - session_id=self.conversation_session_id, - task_id=self.task_id, - iteration_number=current_iteration, - prompt=prompt, - model=self.model, - temperature=0.3, - metadata={ - "task_description": ( - task[:200] + "..." if len(task) > 200 else task - ), - "language": language, - "include_tests": include_tests, - "include_docs": include_docs, - "file_path": file_path, - "workspace_path": self.workspace_path, - }, - ) - - # Record start time for response time tracking - start_time = time.time() - - # Call Claude API - response = self.client.messages.create( - model=self.model, - max_tokens=4096, - temperature=0.3, # Lower temperature for more consistent code - messages=[{"role": "user", "content": prompt}], - ) - - # Extract response content and token usage - response_content = response.content[0].text - token_count = ( - getattr(response.usage, "output_tokens", None) - if hasattr(response, "usage") - else None - ) - - # Store Claude response in conversation history - if self.conversation_session_id and self.task_id: - await self.conversation_manager.store_claude_response( - session_id=self.conversation_session_id, - task_id=self.task_id, - iteration_number=current_iteration, - response=response_content, - token_count=token_count, - model=self.model, - start_time=start_time, - metadata={ - "prompt_length": len(prompt), - "response_length": len(response_content), - }, - ) - - # Parse the response and extract code files - result = self._parse_response( - response_content, language, include_tests, include_docs - ) - - # Save files to workspace if available - saved_files = [] - if self.workspace_path and os.path.exists(self.workspace_path): - saved_files = self._save_files_to_workspace(result["files"]) - - # Store code generations in database - if self.task_id: - for file_info in result["files"]: - await self.conversation_manager.store_code_generation( - task_id=self.task_id, - iteration_number=current_iteration, - file_path=file_info["filename"], - file_type=file_info["type"], - language=file_info.get("language", language), - content=file_info["content"], - ) - - # Update repository context with changed files - self.repository_context["files_changed"].extend( - [ - f["filename"] - for f in result["files"] - if f["type"] == "implementation" - ] - ) - - # Run tests if generated and in workspace - test_results = None - if include_tests and any(f["type"] == "test" for f in result["files"]): - if self.workspace_path and os.path.exists(self.workspace_path): - test_results = self._run_tests(self.workspace_path, language) - - # Store test results - if self.conversation_session_id and self.task_id: - await self.conversation_manager.store_message( - session_id=self.conversation_session_id, - message={ - "task_id": self.task_id, - "iteration_number": current_iteration, - "message_type": MessageType.TEST_RESULT, - "content": json.dumps(test_results), - "metadata": { - "test_framework": ( - "pytest" if language == "python" else "jest" - ), - "workspace_path": self.workspace_path, - }, - }, - ) - - return { - "status": "success", - "files": result["files"], - "saved_files": saved_files, - "explanation": result.get("explanation", ""), - "test_results": test_results, - "commit_message": result.get("commit_message", ""), - "execution_summary": f"Generated {len(result['files'])} files for {language} task: {task[:100]}...", - "iteration": current_iteration, - "workspace_path": self.workspace_path, - "repository_context": self.repository_context, - "conversation_tracked": bool(self.conversation_session_id), - "token_count": token_count, - } - - except Exception as e: - # Store error in conversation history - if self.conversation_session_id and self.task_id: - try: - await self.conversation_manager.store_message( - session_id=self.conversation_session_id, - message={ - "task_id": self.task_id, - "iteration_number": self.repository_context[ - "iteration_count" - ], - "message_type": MessageType.ERROR_MESSAGE, - "content": str(e), - "metadata": { - "error_type": type(e).__name__, - "task_description": ( - task[:200] + "..." if len(task) > 200 else task - ), - }, - }, - ) - except Exception as conv_error: - # Don't let conversation storage errors break the main flow - print( - f"Warning: Failed to store error in conversation: {conv_error}" - ) - - return { - "status": "error", - "error": str(e), - "error_type": "execution_error", - "iteration": self.repository_context["iteration_count"], - "conversation_tracked": bool(self.conversation_session_id), - } +import asyncio +import json +import os +import subprocess # nosec B404 -- used only with static executable names + arg lists and shell=False (see _run_tests) +import tempfile +import time +from pathlib import Path +from typing import Any, Dict, List, Optional, Type + +# Import Anthropic SDK for real Claude integration +import anthropic +from anthropic import Anthropic +from crewai.tools import BaseTool +from pydantic import BaseModel, Field + +# Import conversation manager for full chat tracking. +# This module is imported both as part of the `services.orchestrator` package +# (relative form, e.g. from main.py/agent_manager.py) and flat with +# services/orchestrator on sys.path (e.g. from tests and main_with_hierarchy.py), +# so support both — mirrors the existing pattern in hierarchy_endpoints.py. +try: + from .conversation_manager import ConversationManager, MessageType +except ImportError: # pragma: no cover - flat import (no parent package) + from conversation_manager import ConversationManager, MessageType + + +class ClaudeCodeInput(BaseModel): + """Input schema for Claude Code tool""" + + task: str = Field(description="Coding task to complete") + language: str = Field(default="python", description="Programming language") + context: str = Field(default="", description="Additional context or requirements") + include_tests: bool = Field( + default=True, description="Whether to include unit tests" + ) + include_docs: bool = Field( + default=True, description="Whether to include documentation" + ) + file_path: Optional[str] = Field( + default=None, description="Optional file path for code context" + ) + + +class ClaudeCodeWrapper(BaseTool): + name: str = "claude_code" + description: str = """ + Execute advanced coding tasks using Claude AI with real-time code generation, + testing, and documentation. Supports multiple programming languages and + follows industry best practices. Enhanced for repository context and Git integration. + """ + args_schema: Type[BaseModel] = ClaudeCodeInput + + # Runtime attributes. ``BaseTool`` is a Pydantic v2 model, which rejects + # assignment to undeclared attributes ("object has no field ..."). These + # are declared as model fields (rather than PrivateAttr) so they remain + # publicly readable on the instance (e.g. ``wrapper.client`` / + # ``wrapper.model``), preserving the tool's public interface. Object + # handles (Anthropic SDK client, git/conversation managers) are typed + # ``Any`` so Pydantic stores them as-is without schema validation. + client: Optional[Any] = None + model: str = "claude-3-5-sonnet-20241022" + workspace_path: Optional[str] = None + git_manager: Optional[Any] = None + agent_id: Optional[str] = None + task_id: Optional[str] = None + conversation_manager: Optional[Any] = None + conversation_session_id: Optional[str] = None + current_context: Dict[str, Any] = Field(default_factory=dict) + repository_context: Dict[str, Any] = Field(default_factory=dict) + + def __init__( + self, + workspace_path: Optional[str] = None, + git_manager: Optional[Any] = None, + agent_id: Optional[str] = None, + task_id: Optional[str] = None, + conversation_manager: Optional[ConversationManager] = None, + ): + super().__init__() + self.client = Anthropic( + api_key=os.getenv("ANTHROPIC_API_KEY"), + ) + self.model = "claude-3-5-sonnet-20241022" + self.workspace_path = workspace_path or os.getcwd() + self.git_manager = git_manager + self.agent_id = agent_id + self.task_id = task_id + self.conversation_manager = conversation_manager or ConversationManager() + self.conversation_session_id: Optional[str] = None + self.current_context = {} # Store context between iterations + + # Repository context + self.repository_context = { + "files_changed": [], + "current_branch": None, + "last_commit": None, + "iteration_count": 0, + } + + def _run( + self, + task: str, + language: str = "python", + context: str = "", + include_tests: bool = True, + include_docs: bool = True, + file_path: Optional[str] = None, + iteration_number: Optional[int] = None, + ) -> str: + """Execute Claude Code for a specific task with real AI integration""" + + try: + # Update iteration count + if iteration_number: + self.repository_context["iteration_count"] = iteration_number + else: + self.repository_context["iteration_count"] += 1 + + # Get repository context if Git manager is available + repo_context = "" + if self.git_manager: + try: + # Note: In a full implementation, we'd make this method async + # For now, we'll skip the Git context in the sync version + repo_context = ( + "Repository context: Available (Git manager configured)" + ) + except Exception as e: + repo_context = f"Repository context unavailable: {str(e)}" + + # Read existing file context if provided + existing_code = "" + if file_path: + # Use workspace-relative path if available + full_path = ( + os.path.join(self.workspace_path, file_path) + if not os.path.isabs(file_path) + else file_path + ) + if os.path.exists(full_path): + with open(full_path, "r") as f: + existing_code = f.read() + + # Prepare the comprehensive prompt with repository context + prompt = self._build_prompt( + task=task, + language=language, + context=context, + existing_code=existing_code, + include_tests=include_tests, + include_docs=include_docs, + repo_context=repo_context, + ) + + # Call Claude API + response = self.client.messages.create( + model=self.model, + max_tokens=4096, + temperature=0.3, # Lower temperature for more consistent code + messages=[{"role": "user", "content": prompt}], + ) + + # Parse the response and extract code files + result = self._parse_response( + response.content[0].text, language, include_tests, include_docs + ) + + # Save files to workspace if available, otherwise use temp directory + if self.workspace_path and os.path.exists(self.workspace_path): + saved_files = self._save_files_to_workspace(result["files"]) + else: + with tempfile.TemporaryDirectory() as tmpdir: + saved_files = self._save_files(result["files"], tmpdir) + + # Update repository context with changed files + self.repository_context["files_changed"].extend( + [ + f["filename"] + for f in result["files"] + if f["type"] == "implementation" + ] + ) + + # Run tests if generated and in workspace + test_results = None + if include_tests and any(f["type"] == "test" for f in result["files"]): + if self.workspace_path and os.path.exists(self.workspace_path): + test_results = self._run_tests(self.workspace_path, language) + else: + with tempfile.TemporaryDirectory() as tmpdir: + self._save_files(result["files"], tmpdir) + test_results = self._run_tests(tmpdir, language) + + return json.dumps( + { + "status": "success", + "files": result["files"], + "explanation": result.get("explanation", ""), + "test_results": test_results, + "commit_message": result.get("commit_message", ""), + "execution_summary": f"Generated {len(result['files'])} files for {language} task: {task[:100]}...", + "iteration": self.repository_context["iteration_count"], + "workspace_path": self.workspace_path, + "repository_context": self.repository_context, + } + ) + + except anthropic.APIError as e: + return json.dumps( + { + "status": "error", + "error": f"Claude API error: {str(e)}", + "error_type": "api_error", + } + ) + except Exception as e: + return json.dumps( + { + "status": "error", + "error": f"Unexpected error: {str(e)}", + "error_type": "general_error", + } + ) + + def _build_prompt( + self, + task: str, + language: str, + context: str, + existing_code: str, + include_tests: bool, + include_docs: bool, + repo_context: str = "", + ) -> str: + """Build a comprehensive prompt for Claude with repository context""" + + # Build agent context + agent_info = "" + if self.agent_id and self.task_id: + agent_info = f""" +**Agent Context**: +- Agent ID: {self.agent_id} +- Task ID: {self.task_id} +- Iteration: {self.repository_context['iteration_count']} +- Workspace: {self.workspace_path} +""" + + prompt = f""" +You are an expert {language} developer working autonomously as part of FuzeAgent AI team. I need you to complete the following coding task: + +{agent_info} + +**Task**: {task} + +**Programming Language**: {language} + +**Additional Context**: {context} + +{repo_context} + +**Existing Code** (if any): +```{language} +{existing_code} +``` + +**Requirements**: +1. Write clean, maintainable, and well-documented code +2. Follow {language} best practices and conventions +3. Include proper error handling +4. Use type hints (where applicable) +5. {"Include comprehensive unit tests" if include_tests else "Focus only on implementation"} +6. {"Include docstrings and comments" if include_docs else "Minimal documentation"} +7. Consider the repository context and maintain consistency with existing code +8. Write code that integrates well with the current branch and recent changes + +**Output Format**: +Please structure your response as follows: + +## Explanation +Brief explanation of your approach and key decisions, considering the repository context. + +## Implementation + +### Main Code +```{language} +# Your main implementation here +``` + +{"### Tests" if include_tests else ""} +{f"```{language}" if include_tests else ""} +{"# Your test code here" if include_tests else ""} +{f"```" if include_tests else ""} + +{"### Documentation" if include_docs else ""} +{"```markdown" if include_docs else ""} +{"# Your documentation here" if include_docs else ""} +{f"```" if include_docs else ""} + +## Commit Message +Suggest a concise git commit message for these changes that follows the repository's commit history style. + +Please ensure the code is production-ready and follows industry standards. +""" + return prompt + + def _parse_response( + self, response: str, language: str, include_tests: bool, include_docs: bool + ) -> Dict[str, Any]: + """Parse Claude's response and extract code files""" + + files = [] + explanation = "" + commit_message = "" + + # Extract explanation + if "## Explanation" in response: + explanation_start = response.find("## Explanation") + len("## Explanation") + explanation_end = response.find("## Implementation") + if explanation_end > explanation_start: + explanation = response[explanation_start:explanation_end].strip() + + # Extract commit message + if "## Commit Message" in response: + commit_start = response.find("## Commit Message") + len("## Commit Message") + commit_message = response[commit_start:].strip() + # Clean up the commit message + commit_message = commit_message.split("\n")[0].strip() + + # Extract main code + main_code = self._extract_code_block(response, "### Main Code", language) + if main_code: + file_ext = self._get_file_extension(language) + files.append( + { + "filename": f"main.{file_ext}", + "content": main_code, + "type": "implementation", + "language": language, + } + ) + + # Extract tests if requested + if include_tests: + test_code = self._extract_code_block(response, "### Tests", language) + if test_code: + test_ext = self._get_file_extension(language) + files.append( + { + "filename": f"test_main.{test_ext}", + "content": test_code, + "type": "test", + "language": language, + } + ) + + # Extract documentation if requested + if include_docs: + docs = self._extract_code_block(response, "### Documentation", "markdown") + if docs: + files.append( + { + "filename": "README.md", + "content": docs, + "type": "documentation", + "language": "markdown", + } + ) + + return { + "files": files, + "explanation": explanation, + "commit_message": commit_message, + } + + def _extract_code_block( + self, text: str, section: str, language: str + ) -> Optional[str]: + """Extract code block from a specific section""" + + section_start = text.find(section) + if section_start == -1: + return None + + # Find the start of the code block + code_start = text.find(f"```{language}", section_start) + if code_start == -1: + code_start = text.find("```", section_start) + if code_start == -1: + return None + + # Find the end of the code block + code_content_start = text.find("\n", code_start) + 1 + code_end = text.find("```", code_content_start) + + if code_end == -1: + return None + + return text[code_content_start:code_end].strip() + + def _get_file_extension(self, language: str) -> str: + """Get appropriate file extension for language""" + extensions = { + "python": "py", + "javascript": "js", + "typescript": "ts", + "java": "java", + "cpp": "cpp", + "c": "c", + "rust": "rs", + "go": "go", + "ruby": "rb", + "php": "php", + "swift": "swift", + "kotlin": "kt", + "scala": "scala", + "r": "R", + "sql": "sql", + "html": "html", + "css": "css", + "shell": "sh", + "bash": "sh", + } + return extensions.get(language.lower(), "txt") + + def _save_files(self, files: List[Dict], tmpdir: str) -> List[str]: + """Save generated files to temporary directory""" + saved_files = [] + + for file_info in files: + file_path = os.path.join(tmpdir, file_info["filename"]) + with open(file_path, "w") as f: + f.write(file_info["content"]) + saved_files.append(file_path) + + return saved_files + + def _run_tests(self, tmpdir: str, language: str) -> Optional[Dict[str, Any]]: + """Run tests for the generated code""" + + try: + if language == "python": + # Try to run pytest + result = subprocess.run( # nosec B603 B607 -- shell=False, static argv; runs pytest on generated code inside an isolated tmpdir + ["python", "-m", "pytest", tmpdir, "-v"], + capture_output=True, + text=True, + timeout=60, + cwd=tmpdir, + ) + + return { + "exit_code": result.returncode, + "stdout": result.stdout, + "stderr": result.stderr, + "success": result.returncode == 0, + } + elif language == "javascript": + # Try to run with node + test_files = [f for f in os.listdir(tmpdir) if f.startswith("test_")] + if test_files: + result = subprocess.run( # nosec B603 B607 -- shell=False, static argv; runs a generated test file inside an isolated tmpdir + ["node", test_files[0]], + capture_output=True, + text=True, + timeout=60, + cwd=tmpdir, + ) + + return { + "exit_code": result.returncode, + "stdout": result.stdout, + "stderr": result.stderr, + "success": result.returncode == 0, + } + + return None + + except subprocess.TimeoutExpired: + return { + "exit_code": -1, + "stdout": "", + "stderr": "Test execution timed out", + "success": False, + } + except Exception as e: + return { + "exit_code": -1, + "stdout": "", + "stderr": f"Test execution error: {str(e)}", + "success": False, + } + + def _build_repository_context( + self, branch_status: Dict[str, Any], commit_history: List[Any] + ) -> str: + """Build repository context string for the prompt""" + context_parts = ["**Repository Context**:"] + + if branch_status: + current_branch = branch_status.get("current_branch", "unknown") + feature_branch = branch_status.get("feature_branch") + has_changes = branch_status.get("has_uncommitted_changes", False) + remote_status = branch_status.get("remote_status", "unknown") + + context_parts.append(f"- Current Branch: `{current_branch}`") + if feature_branch: + context_parts.append(f"- Feature Branch: `{feature_branch}`") + context_parts.append( + f"- Uncommitted Changes: {'Yes' if has_changes else 'No'}" + ) + context_parts.append(f"- Remote Status: {remote_status}") + + if commit_history: + context_parts.append("- Recent Commits:") + for i, commit in enumerate(commit_history[:3]): + context_parts.append(f" {i+1}. `{commit.hash[:8]}` - {commit.message}") + if commit.files_changed: + context_parts.append( + f" Files: {', '.join(commit.files_changed[:5])}" + ) + + if self.repository_context.get("files_changed"): + changed_files = list(set(self.repository_context["files_changed"])) + context_parts.append( + f"- Files Modified This Session: {', '.join(changed_files)}" + ) + + return "\n".join(context_parts) + "\n" + + def _save_files_to_workspace(self, files: List[Dict]) -> List[str]: + """Save generated files directly to workspace""" + saved_files = [] + + for file_info in files: + file_path = os.path.join(self.workspace_path, file_info["filename"]) + + # Create directory if needed + os.makedirs(os.path.dirname(file_path), exist_ok=True) + + with open(file_path, "w") as f: + f.write(file_info["content"]) + saved_files.append(file_path) + + return saved_files + + async def commit_and_push_changes( + self, commit_message: str, files: Optional[List[str]] = None + ) -> Dict[str, Any]: + """Commit and push changes using Git manager""" + if not self.git_manager: + return {"success": False, "error": "No Git manager available"} + + try: + # Commit changes + commit_hash = await self.git_manager.commit_changes( + message=commit_message, + files=files, + iteration_number=self.repository_context["iteration_count"], + ) + + if commit_hash: + self.repository_context["last_commit"] = commit_message + return { + "success": True, + "commit_hash": commit_hash, + "message": "Changes committed successfully", + } + else: + return {"success": True, "message": "No changes to commit"} + + except Exception as e: + return {"success": False, "error": f"Failed to commit changes: {str(e)}"} + + def get_repository_context(self) -> Dict[str, Any]: + """Get current repository context""" + return self.repository_context.copy() + + def reset_context(self): + """Reset the repository context""" + self.repository_context = { + "files_changed": [], + "current_branch": None, + "last_commit": None, + "iteration_count": 0, + } + + async def start_conversation_session(self, sandbox_id: str) -> str: + """Start a conversation session for tracking all Claude Code interactions""" + if not self.agent_id or not self.task_id: + raise ValueError("Agent ID and Task ID required for conversation tracking") + + self.conversation_session_id = ( + await self.conversation_manager.start_conversation_session( + agent_id=self.agent_id, + task_id=self.task_id, + sandbox_id=sandbox_id, + metadata={ + "workspace_path": self.workspace_path, + "model": self.model, + "git_enabled": bool(self.git_manager), + }, + ) + ) + return self.conversation_session_id + + async def end_conversation_session(self) -> bool: + """End the current conversation session""" + if not self.conversation_session_id: + return False + + success = await self.conversation_manager.end_conversation_session( + self.conversation_session_id + ) + self.conversation_session_id = None + return success + + async def execute_task_async( + self, + task: str, + language: str = "python", + context: str = "", + include_tests: bool = True, + include_docs: bool = True, + file_path: Optional[str] = None, + iteration_number: Optional[int] = None, + ) -> Dict[str, Any]: + """Async version of task execution with full conversation tracking""" + + try: + # Update iteration count + if iteration_number: + self.repository_context["iteration_count"] = iteration_number + else: + self.repository_context["iteration_count"] += 1 + + current_iteration = self.repository_context["iteration_count"] + + # Get repository context if Git manager is available + repo_context = "" + if self.git_manager: + try: + branch_status = await self.git_manager.get_branch_status() + self.repository_context["current_branch"] = branch_status.get( + "current_branch" + ) + + commit_history = await self.git_manager.get_commit_history(limit=3) + if commit_history: + self.repository_context["last_commit"] = commit_history[ + 0 + ].message + + repo_context = self._build_repository_context( + branch_status, commit_history + ) + except Exception as e: + repo_context = f"Repository context unavailable: {str(e)}" + + # Read existing file context if provided + existing_code = "" + if file_path: + # Use workspace-relative path if available + full_path = ( + os.path.join(self.workspace_path, file_path) + if not os.path.isabs(file_path) + else file_path + ) + if os.path.exists(full_path): + with open(full_path, "r") as f: + existing_code = f.read() + + # Prepare the comprehensive prompt with repository context + prompt = self._build_prompt( + task=task, + language=language, + context=context, + existing_code=existing_code, + include_tests=include_tests, + include_docs=include_docs, + repo_context=repo_context, + ) + + # Store user prompt in conversation history + if self.conversation_session_id and self.task_id: + await self.conversation_manager.store_user_prompt( + session_id=self.conversation_session_id, + task_id=self.task_id, + iteration_number=current_iteration, + prompt=prompt, + model=self.model, + temperature=0.3, + metadata={ + "task_description": ( + task[:200] + "..." if len(task) > 200 else task + ), + "language": language, + "include_tests": include_tests, + "include_docs": include_docs, + "file_path": file_path, + "workspace_path": self.workspace_path, + }, + ) + + # Record start time for response time tracking + start_time = time.time() + + # Call Claude API + response = self.client.messages.create( + model=self.model, + max_tokens=4096, + temperature=0.3, # Lower temperature for more consistent code + messages=[{"role": "user", "content": prompt}], + ) + + # Extract response content and token usage + response_content = response.content[0].text + token_count = ( + getattr(response.usage, "output_tokens", None) + if hasattr(response, "usage") + else None + ) + + # Store Claude response in conversation history + if self.conversation_session_id and self.task_id: + await self.conversation_manager.store_claude_response( + session_id=self.conversation_session_id, + task_id=self.task_id, + iteration_number=current_iteration, + response=response_content, + token_count=token_count, + model=self.model, + start_time=start_time, + metadata={ + "prompt_length": len(prompt), + "response_length": len(response_content), + }, + ) + + # Parse the response and extract code files + result = self._parse_response( + response_content, language, include_tests, include_docs + ) + + # Save files to workspace if available + saved_files = [] + if self.workspace_path and os.path.exists(self.workspace_path): + saved_files = self._save_files_to_workspace(result["files"]) + + # Store code generations in database + if self.task_id: + for file_info in result["files"]: + await self.conversation_manager.store_code_generation( + task_id=self.task_id, + iteration_number=current_iteration, + file_path=file_info["filename"], + file_type=file_info["type"], + language=file_info.get("language", language), + content=file_info["content"], + ) + + # Update repository context with changed files + self.repository_context["files_changed"].extend( + [ + f["filename"] + for f in result["files"] + if f["type"] == "implementation" + ] + ) + + # Run tests if generated and in workspace + test_results = None + if include_tests and any(f["type"] == "test" for f in result["files"]): + if self.workspace_path and os.path.exists(self.workspace_path): + test_results = self._run_tests(self.workspace_path, language) + + # Store test results + if self.conversation_session_id and self.task_id: + await self.conversation_manager.store_message( + session_id=self.conversation_session_id, + message={ + "task_id": self.task_id, + "iteration_number": current_iteration, + "message_type": MessageType.TEST_RESULT, + "content": json.dumps(test_results), + "metadata": { + "test_framework": ( + "pytest" if language == "python" else "jest" + ), + "workspace_path": self.workspace_path, + }, + }, + ) + + return { + "status": "success", + "files": result["files"], + "saved_files": saved_files, + "explanation": result.get("explanation", ""), + "test_results": test_results, + "commit_message": result.get("commit_message", ""), + "execution_summary": f"Generated {len(result['files'])} files for {language} task: {task[:100]}...", + "iteration": current_iteration, + "workspace_path": self.workspace_path, + "repository_context": self.repository_context, + "conversation_tracked": bool(self.conversation_session_id), + "token_count": token_count, + } + + except Exception as e: + # Store error in conversation history + if self.conversation_session_id and self.task_id: + try: + await self.conversation_manager.store_message( + session_id=self.conversation_session_id, + message={ + "task_id": self.task_id, + "iteration_number": self.repository_context[ + "iteration_count" + ], + "message_type": MessageType.ERROR_MESSAGE, + "content": str(e), + "metadata": { + "error_type": type(e).__name__, + "task_description": ( + task[:200] + "..." if len(task) > 200 else task + ), + }, + }, + ) + except Exception as conv_error: + # Don't let conversation storage errors break the main flow + print( + f"Warning: Failed to store error in conversation: {conv_error}" + ) + + return { + "status": "error", + "error": str(e), + "error_type": "execution_error", + "iteration": self.repository_context["iteration_count"], + "conversation_tracked": bool(self.conversation_session_id), + } diff --git a/services/orchestrator/claude_sdk_manager.py b/services/orchestrator/claude_sdk_manager.py index 282da78..3be2939 100644 --- a/services/orchestrator/claude_sdk_manager.py +++ b/services/orchestrator/claude_sdk_manager.py @@ -1,502 +1,502 @@ -""" -Claude SDK Manager for FuzeAgent - -Manages Claude Code SDK processes, handles interactive states, and integrates -with the File Operations Engine to apply code changes safely. -""" - -import asyncio -import json -import logging -import os -import re -import subprocess # nosec B404 -- used with asyncio.create_subprocess_exec (shell=False) and a static arg list -import time -from dataclasses import dataclass -from datetime import datetime -from enum import Enum -from typing import Any, AsyncGenerator, Callable, Dict, List, Optional - -from .conversation_manager import ConversationManager, MessageType -from .file_operations_engine import FileOperationsEngine, OperationBatch - -logger = logging.getLogger(__name__) - - -class ClaudeSDKState(str, Enum): - IDLE = "idle" - INITIALIZING = "initializing" - RUNNING = "running" - WAITING_FOR_INPUT = "waiting_for_input" - WAITING_FOR_APPROVAL = "waiting_for_approval" - PROCESSING = "processing" - ERROR = "error" - COMPLETED = "completed" - TERMINATED = "terminated" - - -class InteractionType(str, Enum): - USER_INPUT = "user_input" - FILE_APPROVAL = "file_approval" - CONFIRMATION = "confirmation" - SELECTION = "selection" - - -@dataclass -class ClaudeInteraction: - """Represents an interaction request from Claude SDK""" - - interaction_id: str - interaction_type: InteractionType - prompt: str - options: Optional[List[str]] = None - default_response: Optional[str] = None - timeout_seconds: Optional[int] = None - metadata: Optional[Dict[str, Any]] = None - - -@dataclass -class ClaudeSDKSession: - """Represents a Claude SDK session""" - - session_id: str - task_id: str - agent_id: str - workspace_path: str - process: Optional[asyncio.subprocess.Process] = None - state: ClaudeSDKState = ClaudeSDKState.IDLE - current_interaction: Optional[ClaudeInteraction] = None - output_buffer: str = "" - error_buffer: str = "" - started_at: Optional[datetime] = None - last_activity: Optional[datetime] = None - - -class ClaudeSDKManager: - """ - Manages Claude Code SDK processes and handles all interactions. - - Features: - - Interactive process management - - Real-time output streaming - - Human-in-the-loop handling - - File operations integration - - State management and recovery - """ - - def __init__( - self, - file_operations_engine: FileOperationsEngine, - conversation_manager: ConversationManager, - ): - self.file_ops_engine = file_operations_engine - self.conversation_manager = conversation_manager - self.sessions: Dict[str, ClaudeSDKSession] = {} - self.interaction_callbacks: Dict[str, Callable] = {} - - # Configuration - self.claude_cli_path = "claude" # Assume in PATH - self.interaction_timeout = 300 # 5 minutes - self.process_timeout = 3600 # 1 hour - - # Pattern matching for interactive states - self.interaction_patterns = { - InteractionType.USER_INPUT: [ - r"Please provide.*?:", - r"Enter your.*?:", - r"What would you like.*?:", - r"\?\s*$", - ], - InteractionType.FILE_APPROVAL: [ - r"Apply these changes.*?\?", - r"Proceed with.*?file.*?changes.*?\?", - r"Create.*?files.*?\?", - r"Modify.*?files.*?\?", - ], - InteractionType.CONFIRMATION: [ - r"Are you sure.*?\?", - r"Continue.*?\?", - r"Proceed.*?\?", - r"\(y/n\)", - ], - InteractionType.SELECTION: [ - r"Choose.*?:", - r"Select.*?:", - r"\[1\].*?\[2\]", - r"Options.*?:", - ], - } - - async def start_session( - self, - task_id: str, - agent_id: str, - workspace_path: str, - task_description: str, - additional_context: Optional[str] = None, - ) -> str: - """Start a new Claude SDK session""" - - session_id = f"claude-{task_id}-{int(time.time())}" - - session = ClaudeSDKSession( - session_id=session_id, - task_id=task_id, - agent_id=agent_id, - workspace_path=workspace_path, - started_at=datetime.now(), - last_activity=datetime.now(), - ) - - self.sessions[session_id] = session - - try: - # Start Claude Code process - await self._start_claude_process( - session, task_description, additional_context - ) - - # Start output monitoring - asyncio.create_task(self._monitor_session(session)) - - logger.info(f"Started Claude SDK session {session_id}") - return session_id - - except Exception as e: - logger.error(f"Error starting Claude SDK session: {e}") - session.state = ClaudeSDKState.ERROR - raise - - async def send_input(self, session_id: str, user_input: str) -> bool: - """Send input to a Claude SDK session""" - - session = self.sessions.get(session_id) - if not session or not session.process: - return False - - try: - # Send input to process - session.process.stdin.write((user_input + "\n").encode()) - await session.process.stdin.drain() - - # Update session state - session.state = ClaudeSDKState.PROCESSING - session.current_interaction = None - session.last_activity = datetime.now() - - # Store interaction in conversation manager - await self.conversation_manager.store_message( - session_id=session_id, - message={ - "task_id": session.task_id, - "iteration_number": 1, # Would be dynamic in real implementation - "message_type": MessageType.USER_PROMPT, - "content": user_input, - "metadata": {"interaction_type": "human_response"}, - }, - ) - - logger.info(f"Sent input to Claude SDK session {session_id}") - return True - - except Exception as e: - logger.error(f"Error sending input to session {session_id}: {e}") - return False - - async def approve_file_operations( - self, session_id: str, batch_id: str, approved: bool - ) -> bool: - """Approve or reject file operations from Claude SDK""" - - # Apply file operations - success = await self.file_ops_engine.approve_operations(batch_id, approved) - - if success and approved: - # Send approval to Claude SDK - await self.send_input(session_id, "y") - return True - elif success and not approved: - # Send rejection to Claude SDK - await self.send_input(session_id, "n") - return True - - return False - - async def get_session_status(self, session_id: str) -> Optional[Dict[str, Any]]: - """Get current status of a Claude SDK session""" - - session = self.sessions.get(session_id) - if not session: - return None - - return { - "session_id": session_id, - "task_id": session.task_id, - "agent_id": session.agent_id, - "state": session.state.value, - "current_interaction": ( - { - "id": session.current_interaction.interaction_id, - "type": session.current_interaction.interaction_type.value, - "prompt": session.current_interaction.prompt, - "options": session.current_interaction.options, - } - if session.current_interaction - else None - ), - "started_at": ( - session.started_at.isoformat() if session.started_at else None - ), - "last_activity": ( - session.last_activity.isoformat() if session.last_activity else None - ), - "workspace_path": session.workspace_path, - } - - async def terminate_session(self, session_id: str) -> bool: - """Terminate a Claude SDK session""" - - session = self.sessions.get(session_id) - if not session: - return False - - try: - if session.process: - session.process.terminate() - try: - await asyncio.wait_for(session.process.wait(), timeout=10) - except asyncio.TimeoutError: - session.process.kill() - await session.process.wait() - - session.state = ClaudeSDKState.TERMINATED - logger.info(f"Terminated Claude SDK session {session_id}") - return True - - except Exception as e: - logger.error(f"Error terminating session {session_id}: {e}") - return False - - def register_interaction_callback(self, session_id: str, callback: Callable): - """Register callback for interaction events""" - self.interaction_callbacks[session_id] = callback - - async def stream_output(self, session_id: str) -> AsyncGenerator[str, None]: - """Stream real-time output from Claude SDK session""" - - session = self.sessions.get(session_id) - if not session: - return - - last_position = 0 - - while session.state not in [ - ClaudeSDKState.COMPLETED, - ClaudeSDKState.TERMINATED, - ClaudeSDKState.ERROR, - ]: - # Check for new output - if len(session.output_buffer) > last_position: - new_output = session.output_buffer[last_position:] - last_position = len(session.output_buffer) - yield new_output - - await asyncio.sleep(0.1) # Small delay to prevent excessive CPU usage - - # Private methods - - async def _start_claude_process( - self, - session: ClaudeSDKSession, - task_description: str, - additional_context: Optional[str] = None, - ): - """Start the Claude Code CLI process""" - - # Build Claude command - cmd = [ - self.claude_cli_path, - "code", - "--workspace", - session.workspace_path, - "--task", - task_description, - ] - - if additional_context: - cmd.extend(["--context", additional_context]) - - # Set environment - env = os.environ.copy() - env["ANTHROPIC_API_KEY"] = os.environ.get("ANTHROPIC_API_KEY", "") - - # Start process - session.process = await asyncio.create_subprocess_exec( - *cmd, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=session.workspace_path, - env=env, - ) - - session.state = ClaudeSDKState.RUNNING - logger.info(f"Started Claude CLI process for session {session.session_id}") - - async def _monitor_session(self, session: ClaudeSDKSession): - """Monitor a Claude SDK session for output and interactions""" - - logger.info(f"Monitoring Claude SDK session {session.session_id}") - - try: - while session.process and session.process.returncode is None: - # Read output with timeout - try: - output_data = await asyncio.wait_for( - session.process.stdout.read(1024), timeout=0.1 - ) - - if output_data: - output_text = output_data.decode("utf-8", errors="replace") - session.output_buffer += output_text - session.last_activity = datetime.now() - - # Process output for interactions - await self._process_output(session, output_text) - - except asyncio.TimeoutError: - # Check for session timeout - if self._is_session_timed_out(session): - logger.warning(f"Session {session.session_id} timed out") - session.state = ClaudeSDKState.ERROR - await self.terminate_session(session.session_id) - break - - # Small delay to prevent excessive CPU usage - await asyncio.sleep(0.01) - - # Process completed - if session.process and session.process.returncode == 0: - session.state = ClaudeSDKState.COMPLETED - logger.info( - f"Claude SDK session {session.session_id} completed successfully" - ) - else: - session.state = ClaudeSDKState.ERROR - logger.error(f"Claude SDK session {session.session_id} failed") - - except Exception as e: - logger.error(f"Error monitoring session {session.session_id}: {e}") - session.state = ClaudeSDKState.ERROR - - async def _process_output(self, session: ClaudeSDKSession, output_text: str): - """Process output from Claude SDK to detect interactions""" - - # Store output in conversation manager - await self.conversation_manager.store_message( - session_id=session.session_id, - message={ - "task_id": session.task_id, - "iteration_number": 1, # Would be dynamic - "message_type": MessageType.CLAUDE_RESPONSE, - "content": output_text, - "metadata": {"stream_chunk": True}, - }, - ) - - # Check for interaction patterns - interaction = self._detect_interaction(output_text) - if interaction: - session.current_interaction = interaction - session.state = ClaudeSDKState.WAITING_FOR_INPUT - - # Notify callback if registered - callback = self.interaction_callbacks.get(session.session_id) - if callback: - asyncio.create_task(callback(session, interaction)) - - logger.info( - f"Detected interaction in session {session.session_id}: {interaction.interaction_type}" - ) - - # Check for file operations - await self._check_for_file_operations(session, output_text) - - def _detect_interaction(self, output_text: str) -> Optional[ClaudeInteraction]: - """Detect if output contains an interaction request""" - - # Check each interaction type - for interaction_type, patterns in self.interaction_patterns.items(): - for pattern in patterns: - if re.search(pattern, output_text, re.IGNORECASE | re.MULTILINE): - # Extract the prompt (last few lines) - lines = output_text.strip().split("\n") - prompt = "\n".join(lines[-3:]) # Last 3 lines as prompt - - interaction_id = f"interaction-{int(time.time())}" - - return ClaudeInteraction( - interaction_id=interaction_id, - interaction_type=interaction_type, - prompt=prompt, - timeout_seconds=self.interaction_timeout, - ) - - return None - - async def _check_for_file_operations( - self, session: ClaudeSDKSession, output_text: str - ): - """Check if output contains file operation requests""" - - # Look for structured file operations (JSON format) - try: - # Try to find JSON blocks in output - json_blocks = re.findall(r"```json\n(.*?)\n```", output_text, re.DOTALL) - for json_block in json_blocks: - try: - operations_data = json.loads(json_block) - if "operations" in operations_data: - # Process file operations - batch = await self.file_ops_engine.process_claude_response( - operations_data, session.task_id, session.agent_id - ) - - if batch.requires_approval: - session.state = ClaudeSDKState.WAITING_FOR_APPROVAL - session.current_interaction = ClaudeInteraction( - interaction_id=f"approval-{batch.batch_id}", - interaction_type=InteractionType.FILE_APPROVAL, - prompt=f"Approve file operations: {batch.description}", - metadata={"batch_id": batch.batch_id}, - ) - else: - # Auto-approved operations - await self.file_ops_engine.apply_operations_if_approved( - batch.batch_id - ) - - logger.info( - f"Detected file operations in session {session.session_id}" - ) - - except json.JSONDecodeError: - continue - - except Exception as e: - logger.error(f"Error processing file operations: {e}") - - def _is_session_timed_out(self, session: ClaudeSDKSession) -> bool: - """Check if session has timed out""" - - if not session.last_activity: - return False - - timeout_seconds = self.process_timeout - if session.current_interaction: - timeout_seconds = ( - session.current_interaction.timeout_seconds or self.interaction_timeout - ) - - time_since_activity = (datetime.now() - session.last_activity).total_seconds() - return time_since_activity > timeout_seconds +""" +Claude SDK Manager for FuzeAgent + +Manages Claude Code SDK processes, handles interactive states, and integrates +with the File Operations Engine to apply code changes safely. +""" + +import asyncio +import json +import logging +import os +import re +import subprocess # nosec B404 -- used with asyncio.create_subprocess_exec (shell=False) and a static arg list +import time +from dataclasses import dataclass +from datetime import datetime +from enum import Enum +from typing import Any, AsyncGenerator, Callable, Dict, List, Optional + +from .conversation_manager import ConversationManager, MessageType +from .file_operations_engine import FileOperationsEngine, OperationBatch + +logger = logging.getLogger(__name__) + + +class ClaudeSDKState(str, Enum): + IDLE = "idle" + INITIALIZING = "initializing" + RUNNING = "running" + WAITING_FOR_INPUT = "waiting_for_input" + WAITING_FOR_APPROVAL = "waiting_for_approval" + PROCESSING = "processing" + ERROR = "error" + COMPLETED = "completed" + TERMINATED = "terminated" + + +class InteractionType(str, Enum): + USER_INPUT = "user_input" + FILE_APPROVAL = "file_approval" + CONFIRMATION = "confirmation" + SELECTION = "selection" + + +@dataclass +class ClaudeInteraction: + """Represents an interaction request from Claude SDK""" + + interaction_id: str + interaction_type: InteractionType + prompt: str + options: Optional[List[str]] = None + default_response: Optional[str] = None + timeout_seconds: Optional[int] = None + metadata: Optional[Dict[str, Any]] = None + + +@dataclass +class ClaudeSDKSession: + """Represents a Claude SDK session""" + + session_id: str + task_id: str + agent_id: str + workspace_path: str + process: Optional[asyncio.subprocess.Process] = None + state: ClaudeSDKState = ClaudeSDKState.IDLE + current_interaction: Optional[ClaudeInteraction] = None + output_buffer: str = "" + error_buffer: str = "" + started_at: Optional[datetime] = None + last_activity: Optional[datetime] = None + + +class ClaudeSDKManager: + """ + Manages Claude Code SDK processes and handles all interactions. + + Features: + - Interactive process management + - Real-time output streaming + - Human-in-the-loop handling + - File operations integration + - State management and recovery + """ + + def __init__( + self, + file_operations_engine: FileOperationsEngine, + conversation_manager: ConversationManager, + ): + self.file_ops_engine = file_operations_engine + self.conversation_manager = conversation_manager + self.sessions: Dict[str, ClaudeSDKSession] = {} + self.interaction_callbacks: Dict[str, Callable] = {} + + # Configuration + self.claude_cli_path = "claude" # Assume in PATH + self.interaction_timeout = 300 # 5 minutes + self.process_timeout = 3600 # 1 hour + + # Pattern matching for interactive states + self.interaction_patterns = { + InteractionType.USER_INPUT: [ + r"Please provide.*?:", + r"Enter your.*?:", + r"What would you like.*?:", + r"\?\s*$", + ], + InteractionType.FILE_APPROVAL: [ + r"Apply these changes.*?\?", + r"Proceed with.*?file.*?changes.*?\?", + r"Create.*?files.*?\?", + r"Modify.*?files.*?\?", + ], + InteractionType.CONFIRMATION: [ + r"Are you sure.*?\?", + r"Continue.*?\?", + r"Proceed.*?\?", + r"\(y/n\)", + ], + InteractionType.SELECTION: [ + r"Choose.*?:", + r"Select.*?:", + r"\[1\].*?\[2\]", + r"Options.*?:", + ], + } + + async def start_session( + self, + task_id: str, + agent_id: str, + workspace_path: str, + task_description: str, + additional_context: Optional[str] = None, + ) -> str: + """Start a new Claude SDK session""" + + session_id = f"claude-{task_id}-{int(time.time())}" + + session = ClaudeSDKSession( + session_id=session_id, + task_id=task_id, + agent_id=agent_id, + workspace_path=workspace_path, + started_at=datetime.now(), + last_activity=datetime.now(), + ) + + self.sessions[session_id] = session + + try: + # Start Claude Code process + await self._start_claude_process( + session, task_description, additional_context + ) + + # Start output monitoring + asyncio.create_task(self._monitor_session(session)) + + logger.info(f"Started Claude SDK session {session_id}") + return session_id + + except Exception as e: + logger.error(f"Error starting Claude SDK session: {e}") + session.state = ClaudeSDKState.ERROR + raise + + async def send_input(self, session_id: str, user_input: str) -> bool: + """Send input to a Claude SDK session""" + + session = self.sessions.get(session_id) + if not session or not session.process: + return False + + try: + # Send input to process + session.process.stdin.write((user_input + "\n").encode()) + await session.process.stdin.drain() + + # Update session state + session.state = ClaudeSDKState.PROCESSING + session.current_interaction = None + session.last_activity = datetime.now() + + # Store interaction in conversation manager + await self.conversation_manager.store_message( + session_id=session_id, + message={ + "task_id": session.task_id, + "iteration_number": 1, # Would be dynamic in real implementation + "message_type": MessageType.USER_PROMPT, + "content": user_input, + "metadata": {"interaction_type": "human_response"}, + }, + ) + + logger.info(f"Sent input to Claude SDK session {session_id}") + return True + + except Exception as e: + logger.error(f"Error sending input to session {session_id}: {e}") + return False + + async def approve_file_operations( + self, session_id: str, batch_id: str, approved: bool + ) -> bool: + """Approve or reject file operations from Claude SDK""" + + # Apply file operations + success = await self.file_ops_engine.approve_operations(batch_id, approved) + + if success and approved: + # Send approval to Claude SDK + await self.send_input(session_id, "y") + return True + elif success and not approved: + # Send rejection to Claude SDK + await self.send_input(session_id, "n") + return True + + return False + + async def get_session_status(self, session_id: str) -> Optional[Dict[str, Any]]: + """Get current status of a Claude SDK session""" + + session = self.sessions.get(session_id) + if not session: + return None + + return { + "session_id": session_id, + "task_id": session.task_id, + "agent_id": session.agent_id, + "state": session.state.value, + "current_interaction": ( + { + "id": session.current_interaction.interaction_id, + "type": session.current_interaction.interaction_type.value, + "prompt": session.current_interaction.prompt, + "options": session.current_interaction.options, + } + if session.current_interaction + else None + ), + "started_at": ( + session.started_at.isoformat() if session.started_at else None + ), + "last_activity": ( + session.last_activity.isoformat() if session.last_activity else None + ), + "workspace_path": session.workspace_path, + } + + async def terminate_session(self, session_id: str) -> bool: + """Terminate a Claude SDK session""" + + session = self.sessions.get(session_id) + if not session: + return False + + try: + if session.process: + session.process.terminate() + try: + await asyncio.wait_for(session.process.wait(), timeout=10) + except asyncio.TimeoutError: + session.process.kill() + await session.process.wait() + + session.state = ClaudeSDKState.TERMINATED + logger.info(f"Terminated Claude SDK session {session_id}") + return True + + except Exception as e: + logger.error(f"Error terminating session {session_id}: {e}") + return False + + def register_interaction_callback(self, session_id: str, callback: Callable): + """Register callback for interaction events""" + self.interaction_callbacks[session_id] = callback + + async def stream_output(self, session_id: str) -> AsyncGenerator[str, None]: + """Stream real-time output from Claude SDK session""" + + session = self.sessions.get(session_id) + if not session: + return + + last_position = 0 + + while session.state not in [ + ClaudeSDKState.COMPLETED, + ClaudeSDKState.TERMINATED, + ClaudeSDKState.ERROR, + ]: + # Check for new output + if len(session.output_buffer) > last_position: + new_output = session.output_buffer[last_position:] + last_position = len(session.output_buffer) + yield new_output + + await asyncio.sleep(0.1) # Small delay to prevent excessive CPU usage + + # Private methods + + async def _start_claude_process( + self, + session: ClaudeSDKSession, + task_description: str, + additional_context: Optional[str] = None, + ): + """Start the Claude Code CLI process""" + + # Build Claude command + cmd = [ + self.claude_cli_path, + "code", + "--workspace", + session.workspace_path, + "--task", + task_description, + ] + + if additional_context: + cmd.extend(["--context", additional_context]) + + # Set environment + env = os.environ.copy() + env["ANTHROPIC_API_KEY"] = os.environ.get("ANTHROPIC_API_KEY", "") + + # Start process + session.process = await asyncio.create_subprocess_exec( + *cmd, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=session.workspace_path, + env=env, + ) + + session.state = ClaudeSDKState.RUNNING + logger.info(f"Started Claude CLI process for session {session.session_id}") + + async def _monitor_session(self, session: ClaudeSDKSession): + """Monitor a Claude SDK session for output and interactions""" + + logger.info(f"Monitoring Claude SDK session {session.session_id}") + + try: + while session.process and session.process.returncode is None: + # Read output with timeout + try: + output_data = await asyncio.wait_for( + session.process.stdout.read(1024), timeout=0.1 + ) + + if output_data: + output_text = output_data.decode("utf-8", errors="replace") + session.output_buffer += output_text + session.last_activity = datetime.now() + + # Process output for interactions + await self._process_output(session, output_text) + + except asyncio.TimeoutError: + # Check for session timeout + if self._is_session_timed_out(session): + logger.warning(f"Session {session.session_id} timed out") + session.state = ClaudeSDKState.ERROR + await self.terminate_session(session.session_id) + break + + # Small delay to prevent excessive CPU usage + await asyncio.sleep(0.01) + + # Process completed + if session.process and session.process.returncode == 0: + session.state = ClaudeSDKState.COMPLETED + logger.info( + f"Claude SDK session {session.session_id} completed successfully" + ) + else: + session.state = ClaudeSDKState.ERROR + logger.error(f"Claude SDK session {session.session_id} failed") + + except Exception as e: + logger.error(f"Error monitoring session {session.session_id}: {e}") + session.state = ClaudeSDKState.ERROR + + async def _process_output(self, session: ClaudeSDKSession, output_text: str): + """Process output from Claude SDK to detect interactions""" + + # Store output in conversation manager + await self.conversation_manager.store_message( + session_id=session.session_id, + message={ + "task_id": session.task_id, + "iteration_number": 1, # Would be dynamic + "message_type": MessageType.CLAUDE_RESPONSE, + "content": output_text, + "metadata": {"stream_chunk": True}, + }, + ) + + # Check for interaction patterns + interaction = self._detect_interaction(output_text) + if interaction: + session.current_interaction = interaction + session.state = ClaudeSDKState.WAITING_FOR_INPUT + + # Notify callback if registered + callback = self.interaction_callbacks.get(session.session_id) + if callback: + asyncio.create_task(callback(session, interaction)) + + logger.info( + f"Detected interaction in session {session.session_id}: {interaction.interaction_type}" + ) + + # Check for file operations + await self._check_for_file_operations(session, output_text) + + def _detect_interaction(self, output_text: str) -> Optional[ClaudeInteraction]: + """Detect if output contains an interaction request""" + + # Check each interaction type + for interaction_type, patterns in self.interaction_patterns.items(): + for pattern in patterns: + if re.search(pattern, output_text, re.IGNORECASE | re.MULTILINE): + # Extract the prompt (last few lines) + lines = output_text.strip().split("\n") + prompt = "\n".join(lines[-3:]) # Last 3 lines as prompt + + interaction_id = f"interaction-{int(time.time())}" + + return ClaudeInteraction( + interaction_id=interaction_id, + interaction_type=interaction_type, + prompt=prompt, + timeout_seconds=self.interaction_timeout, + ) + + return None + + async def _check_for_file_operations( + self, session: ClaudeSDKSession, output_text: str + ): + """Check if output contains file operation requests""" + + # Look for structured file operations (JSON format) + try: + # Try to find JSON blocks in output + json_blocks = re.findall(r"```json\n(.*?)\n```", output_text, re.DOTALL) + for json_block in json_blocks: + try: + operations_data = json.loads(json_block) + if "operations" in operations_data: + # Process file operations + batch = await self.file_ops_engine.process_claude_response( + operations_data, session.task_id, session.agent_id + ) + + if batch.requires_approval: + session.state = ClaudeSDKState.WAITING_FOR_APPROVAL + session.current_interaction = ClaudeInteraction( + interaction_id=f"approval-{batch.batch_id}", + interaction_type=InteractionType.FILE_APPROVAL, + prompt=f"Approve file operations: {batch.description}", + metadata={"batch_id": batch.batch_id}, + ) + else: + # Auto-approved operations + await self.file_ops_engine.apply_operations_if_approved( + batch.batch_id + ) + + logger.info( + f"Detected file operations in session {session.session_id}" + ) + + except json.JSONDecodeError: + continue + + except Exception as e: + logger.error(f"Error processing file operations: {e}") + + def _is_session_timed_out(self, session: ClaudeSDKSession) -> bool: + """Check if session has timed out""" + + if not session.last_activity: + return False + + timeout_seconds = self.process_timeout + if session.current_interaction: + timeout_seconds = ( + session.current_interaction.timeout_seconds or self.interaction_timeout + ) + + time_since_activity = (datetime.now() - session.last_activity).total_seconds() + return time_since_activity > timeout_seconds diff --git a/services/orchestrator/context_enhancement_service.py b/services/orchestrator/context_enhancement_service.py index 60e2c9c..bc9ae8b 100644 --- a/services/orchestrator/context_enhancement_service.py +++ b/services/orchestrator/context_enhancement_service.py @@ -1,774 +1,774 @@ -""" -Context Enhancement Service for FuzeAgent - -This service enhances agent context with relevant organizational and team knowledge -before task execution. It provides intelligent knowledge injection based on -task type, agent capabilities, and historical success patterns. -""" - -import asyncio -import json -import logging -from dataclasses import dataclass -from datetime import datetime -from typing import Any, Dict, List, Optional, Tuple - -import asyncpg - -from .organization_rag_manager import ( - ContentType, - KnowledgeCategory, - KnowledgeSearchResult, - OrganizationRAGManager, -) -from .team_knowledge_manager import TeamKnowledgeManager, TeamKnowledgeSearchResult - -logger = logging.getLogger(__name__) - - -@dataclass -class ContextEnhancement: - """Represents an enhancement to agent context""" - - knowledge_id: str - title: str - content: str - source_type: str # 'organization', 'team', 'agent' - category: str - relevance_score: float - confidence_score: float - usage_stats: Dict[str, Any] - metadata: Dict[str, Any] - - -@dataclass -class EnhancedContext: - """Enhanced context for agent task execution""" - - task_id: str - agent_id: str - team_id: str - organization_id: str - base_context: Dict[str, Any] - organizational_knowledge: List[ContextEnhancement] - team_knowledge: List[ContextEnhancement] - similar_task_insights: List[ContextEnhancement] - success_patterns: List[str] - common_pitfalls: List[str] - recommended_approaches: List[str] - context_summary: str - enhancement_metadata: Dict[str, Any] - - -class ContextEnhancementService: - """ - Enhances agent context with relevant organizational knowledge - to improve task execution success rates. - """ - - def __init__( - self, - database_url: str, - org_rag_manager: OrganizationRAGManager, - team_knowledge_manager: TeamKnowledgeManager, - ): - self.database_url = database_url - self.org_rag_manager = org_rag_manager - self.team_knowledge_manager = team_knowledge_manager - self.pool: Optional[asyncpg.Pool] = None - - # Configuration - self.max_org_knowledge_items = 5 - self.max_team_knowledge_items = 8 - self.max_similar_tasks = 3 - self.min_relevance_threshold = 0.4 - self.context_freshness_days = 90 - - # Statistics - self.enhancements_created = 0 - self.average_enhancement_score = 0.0 - self.knowledge_usage_tracking = {} - - async def initialize(self): - """Initialize the context enhancement service""" - logger.info("Initializing ContextEnhancementService") - - try: - self.pool = await asyncpg.create_pool( - self.database_url, min_size=1, max_size=5, command_timeout=60 - ) - - logger.info("ContextEnhancementService initialized successfully") - - except Exception as e: - logger.error(f"Failed to initialize ContextEnhancementService: {e}") - raise - - async def close(self): - """Close database connections""" - if self.pool: - await self.pool.close() - logger.info("ContextEnhancementService closed") - - async def enhance_agent_context( - self, - agent_id: str, - task_data: Dict[str, Any], - base_context: Optional[Dict[str, Any]] = None, - ) -> EnhancedContext: - """Enhance agent context with relevant organizational knowledge""" - - try: - # Get agent and team information - agent_info = await self._get_agent_info(agent_id) - if not agent_info: - raise ValueError(f"Agent {agent_id} not found") - - # Build search queries based on task data - search_queries = self._build_search_queries(task_data, agent_info) - - # Gather knowledge from different sources - org_knowledge = await self._gather_organizational_knowledge( - agent_info["organization_id"], - search_queries, - agent_id, - agent_info["team_id"], - ) - - team_knowledge = await self._gather_team_knowledge( - agent_info["team_id"], search_queries - ) - - similar_tasks = await self._find_similar_task_insights( - agent_info["organization_id"], task_data, agent_id - ) - - # Extract patterns and recommendations - success_patterns = await self._extract_success_patterns( - org_knowledge + team_knowledge + similar_tasks - ) - - pitfalls = await self._extract_common_pitfalls( - agent_info["organization_id"], task_data - ) - - recommendations = await self._generate_recommendations( - task_data, org_knowledge, team_knowledge, similar_tasks - ) - - # Create context summary - context_summary = self._create_context_summary( - task_data, org_knowledge, team_knowledge, success_patterns - ) - - # Build enhanced context - enhanced_context = EnhancedContext( - task_id=task_data.get("task_id", ""), - agent_id=agent_id, - team_id=agent_info["team_id"], - organization_id=agent_info["organization_id"], - base_context=base_context or {}, - organizational_knowledge=org_knowledge, - team_knowledge=team_knowledge, - similar_task_insights=similar_tasks, - success_patterns=success_patterns, - common_pitfalls=pitfalls, - recommended_approaches=recommendations, - context_summary=context_summary, - enhancement_metadata={ - "enhancement_timestamp": datetime.now().isoformat(), - "search_queries_used": search_queries, - "knowledge_sources_count": { - "organizational": len(org_knowledge), - "team": len(team_knowledge), - "similar_tasks": len(similar_tasks), - }, - "total_relevance_score": sum( - item.relevance_score - for item in org_knowledge + team_knowledge + similar_tasks - ), - "enhancement_version": "1.0", - }, - ) - - # Track enhancement usage - await self._track_enhancement_usage(enhanced_context) - - self.enhancements_created += 1 - - logger.info( - f"Enhanced context for agent {agent_id}: " - f"{len(org_knowledge)} org + {len(team_knowledge)} team + " - f"{len(similar_tasks)} similar task insights" - ) - - return enhanced_context - - except Exception as e: - logger.error(f"Error enhancing context for agent {agent_id}: {e}") - # Return minimal enhanced context on error - return EnhancedContext( - task_id=task_data.get("task_id", ""), - agent_id=agent_id, - team_id="", - organization_id="", - base_context=base_context or {}, - organizational_knowledge=[], - team_knowledge=[], - similar_task_insights=[], - success_patterns=[], - common_pitfalls=[], - recommended_approaches=[], - context_summary="Context enhancement failed - using minimal context", - enhancement_metadata={"error": str(e)}, - ) - - async def get_contextual_guidance( - self, - agent_id: str, - current_task_context: Dict[str, Any], - current_iteration: int = 1, - ) -> Dict[str, Any]: - """Get contextual guidance during task execution""" - - try: - agent_info = await self._get_agent_info(agent_id) - if not agent_info: - return {"guidance": [], "suggestions": []} - - # Build guidance based on current context - guidance_items = [] - - # Get iteration-specific guidance - if current_iteration > 3: - guidance_items.extend( - await self._get_iteration_guidance( - agent_info["organization_id"], current_iteration - ) - ) - - # Get context-specific suggestions - suggestions = await self._get_contextual_suggestions( - agent_info["organization_id"], - agent_info["team_id"], - current_task_context, - ) - - return { - "guidance": guidance_items, - "suggestions": suggestions, - "iteration": current_iteration, - "generated_at": datetime.now().isoformat(), - } - - except Exception as e: - logger.error(f"Error getting contextual guidance: {e}") - return {"guidance": [], "suggestions": []} - - async def update_knowledge_effectiveness( - self, - knowledge_id: str, - knowledge_source: str, - task_success: bool, - agent_feedback: Optional[Dict[str, Any]] = None, - ): - """Update knowledge effectiveness based on usage outcomes""" - - try: - if knowledge_source == "organization": - await self.org_rag_manager.update_knowledge_quality( - knowledge_id=knowledge_id, - success_correlation=1.0 if task_success else -0.2, - feedback_metadata=agent_feedback, - ) - elif knowledge_source == "team": - # Update team knowledge effectiveness - async with self.pool.acquire() as conn: - agent_id = ( - agent_feedback.get("agent_id") if agent_feedback else None - ) - if agent_id: - await self.team_knowledge_manager.update_knowledge_effectiveness( - team_knowledge_id=knowledge_id, - agent_id=agent_id, - task_success=task_success, - feedback_score=agent_feedback.get("usefulness_score"), - usage_context=agent_feedback, - ) - - # Track in local usage statistics - if knowledge_id not in self.knowledge_usage_tracking: - self.knowledge_usage_tracking[knowledge_id] = { - "usage_count": 0, - "success_count": 0, - "effectiveness_score": 0.0, - } - - stats = self.knowledge_usage_tracking[knowledge_id] - stats["usage_count"] += 1 - if task_success: - stats["success_count"] += 1 - stats["effectiveness_score"] = stats["success_count"] / stats["usage_count"] - - except Exception as e: - logger.error(f"Error updating knowledge effectiveness: {e}") - - async def get_enhancement_statistics( - self, - organization_id: Optional[str] = None, - team_id: Optional[str] = None, - days_back: int = 30, - ) -> Dict[str, Any]: - """Get context enhancement statistics""" - - try: - async with self.pool.acquire() as conn: - # Basic enhancement statistics - where_conditions = [ - "created_at >= NOW() - INTERVAL '%s days'" % days_back - ] - params = [] - - if organization_id: - where_conditions.append("organization_id = $1") - params.append(organization_id) - - if team_id: - where_conditions.append( - "team_id = $2" if organization_id else "team_id = $1" - ) - params.append(team_id) - - # This is a placeholder - in practice you'd have a table to track enhancements - stats = { - "total_enhancements": self.enhancements_created, - "average_knowledge_items_per_enhancement": { - "organizational": 3.2, - "team": 4.1, - "similar_tasks": 1.8, - }, - "knowledge_effectiveness": dict(self.knowledge_usage_tracking), - "generated_at": datetime.now().isoformat(), - } - - return stats - - except Exception as e: - logger.error(f"Error getting enhancement statistics: {e}") - return {} - - async def _get_agent_info(self, agent_id: str) -> Optional[Dict[str, Any]]: - """Get agent information including team and organization""" - - async with self.pool.acquire() as conn: - agent_info = await conn.fetchrow( - """ - SELECT a.id, a.name, a.type, a.config, a.team_id, - t.organization_id, t.name as team_name, - o.name as organization_name - FROM agents a - JOIN teams t ON a.team_id = t.id - JOIN organizations o ON t.organization_id = o.id - WHERE a.id = $1 - """, - agent_id, - ) - - return dict(agent_info) if agent_info else None - - def _build_search_queries( - self, task_data: Dict[str, Any], agent_info: Dict[str, Any] - ) -> List[str]: - """Build search queries based on task data and agent information""" - - queries = [] - - # Primary query from task description - if task_data.get("description"): - queries.append(task_data["description"][:200]) - - # Query from task title - if task_data.get("title"): - queries.append(task_data["title"]) - - # Technology-specific queries - if task_data.get("technologies"): - for tech in task_data["technologies"]: - queries.append(f"{tech} development best practices") - - # Task type specific query - if task_data.get("task_type"): - queries.append(f"{task_data['task_type']} implementation guide") - - # Agent type specific query - agent_type = agent_info.get("type", "") - if agent_type: - queries.append(f"{agent_type} workflow best practices") - - return queries[:5] # Limit to top 5 queries - - async def _gather_organizational_knowledge( - self, - organization_id: str, - search_queries: List[str], - agent_id: str, - team_id: str, - ) -> List[ContextEnhancement]: - """Gather relevant organizational knowledge""" - - org_knowledge = [] - - for query in search_queries: - search_results = await self.org_rag_manager.search_knowledge( - organization_id=organization_id, - query=query, - limit=self.max_org_knowledge_items // len(search_queries) + 1, - min_similarity=self.min_relevance_threshold, - requester_agent_id=agent_id, - requester_team_id=team_id, - ) - - for result in search_results: - if result.combined_score >= self.min_relevance_threshold: - enhancement = ContextEnhancement( - knowledge_id=result.knowledge.id, - title=result.knowledge.title, - content=( - result.knowledge.content[:1000] + "..." - if len(result.knowledge.content) > 1000 - else result.knowledge.content - ), - source_type="organization", - category=result.knowledge.knowledge_category.value, - relevance_score=result.combined_score, - confidence_score=result.knowledge.quality_score, - usage_stats={ - "usage_count": result.knowledge.usage_count, - "success_correlation": result.knowledge.success_correlation, - }, - metadata=result.knowledge.metadata, - ) - org_knowledge.append(enhancement) - - # Sort by relevance and remove duplicates - seen_ids = set() - unique_knowledge = [] - for item in sorted( - org_knowledge, key=lambda x: x.relevance_score, reverse=True - ): - if item.knowledge_id not in seen_ids: - unique_knowledge.append(item) - seen_ids.add(item.knowledge_id) - - return unique_knowledge[: self.max_org_knowledge_items] - - async def _gather_team_knowledge( - self, team_id: str, search_queries: List[str] - ) -> List[ContextEnhancement]: - """Gather relevant team knowledge""" - - team_knowledge = [] - - for query in search_queries: - search_results = await self.team_knowledge_manager.search_team_knowledge( - team_id=team_id, - query=query, - limit=self.max_team_knowledge_items // len(search_queries) + 1, - min_similarity=self.min_relevance_threshold, - ) - - for result in search_results: - if result.combined_score >= self.min_relevance_threshold: - enhancement = ContextEnhancement( - knowledge_id=result.team_knowledge.id, - title=result.team_knowledge.title, - content=( - result.team_knowledge.content[:1000] + "..." - if len(result.team_knowledge.content) > 1000 - else result.team_knowledge.content - ), - source_type="team", - category=result.team_knowledge.knowledge_category.value, - relevance_score=result.combined_score, - confidence_score=result.team_knowledge.effectiveness_score, - usage_stats={ - "adoption_rate": result.team_knowledge.agent_adoption_rate, - "effectiveness": result.team_knowledge.effectiveness_score, - }, - metadata=result.team_knowledge.metadata, - ) - team_knowledge.append(enhancement) - - # Sort and deduplicate - seen_ids = set() - unique_knowledge = [] - for item in sorted( - team_knowledge, key=lambda x: x.relevance_score, reverse=True - ): - if item.knowledge_id not in seen_ids: - unique_knowledge.append(item) - seen_ids.add(item.knowledge_id) - - return unique_knowledge[: self.max_team_knowledge_items] - - async def _find_similar_task_insights( - self, organization_id: str, task_data: Dict[str, Any], agent_id: str - ) -> List[ContextEnhancement]: - """Find insights from similar completed tasks""" - - similar_tasks = [] - - try: - async with self.pool.acquire() as conn: - # Find similar tasks based on description similarity and success - similar_task_data = await conn.fetch( - """ - SELECT t.id, t.title, t.description, t.result, t.completed_at, - a.name as agent_name, a.type as agent_type, - similarity(t.description, $2) as similarity_score - FROM tasks t - JOIN agents a ON t.agent_id = a.id - JOIN teams te ON a.team_id = te.id - WHERE te.organization_id = $1 - AND t.status = 'completed' - AND t.result->>'status' = 'completed' - AND t.completed_at >= NOW() - INTERVAL '90 days' - AND similarity(t.description, $2) > 0.3 - ORDER BY similarity_score DESC, t.completed_at DESC - LIMIT $3 - """, - organization_id, - task_data.get("description", ""), - self.max_similar_tasks, - ) - - for task in similar_task_data: - # Extract insights from the task result - task_result = ( - task["result"] if isinstance(task["result"], dict) else {} - ) - - insights_content = self._extract_task_insights( - dict(task), task_result - ) - - if insights_content: - enhancement = ContextEnhancement( - knowledge_id=str(task["id"]), - title=f"Similar Task: {task['title'][:50]}...", - content=insights_content, - source_type="similar_task", - category="process", - relevance_score=float(task["similarity_score"]), - confidence_score=0.8, # High confidence for successful completed tasks - usage_stats={ - "agent_type": task["agent_type"], - "completion_date": task["completed_at"].isoformat(), - }, - metadata={ - "source_task_id": str(task["id"]), - "source_agent": task["agent_name"], - "similarity_score": float(task["similarity_score"]), - }, - ) - similar_tasks.append(enhancement) - - except Exception as e: - logger.error(f"Error finding similar task insights: {e}") - - return similar_tasks - - async def _extract_success_patterns( - self, all_knowledge: List[ContextEnhancement] - ) -> List[str]: - """Extract success patterns from knowledge items""" - - patterns = [] - - for item in all_knowledge: - # Look for success indicators in metadata - if "success_indicators" in item.metadata: - patterns.extend(item.metadata["success_indicators"]) - - # Extract patterns from high-confidence, high-usage items - if item.confidence_score > 0.7 and item.relevance_score > 0.6: - if "optimization" in item.title.lower(): - patterns.append("Focus on optimization early") - if "test" in item.title.lower(): - patterns.append("Comprehensive testing leads to success") - if "pattern" in item.title.lower(): - patterns.append("Follow established patterns") - - return list(set(patterns)) # Remove duplicates - - async def _extract_common_pitfalls( - self, organization_id: str, task_data: Dict[str, Any] - ) -> List[str]: - """Extract common pitfalls for this type of task""" - - pitfalls = [] - - try: - # Search for error patterns and failure knowledge - error_knowledge = await self.org_rag_manager.search_knowledge( - organization_id=organization_id, - query=f"error pattern {task_data.get('task_type', '')}", - categories=[KnowledgeCategory.TROUBLESHOOTING], - limit=5, - min_similarity=0.3, - ) - - for result in error_knowledge: - if "failure_patterns" in result.knowledge.metadata: - pitfalls.extend(result.knowledge.metadata["failure_patterns"]) - - # Extract pitfalls from error pattern content - content_lower = result.knowledge.content.lower() - if ( - "avoid" in content_lower - or "pitfall" in content_lower - or "common mistake" in content_lower - ): - pitfalls.append(result.knowledge.title) - - except Exception as e: - logger.error(f"Error extracting pitfalls: {e}") - - return list(set(pitfalls))[:5] # Top 5 pitfalls - - async def _generate_recommendations( - self, - task_data: Dict[str, Any], - org_knowledge: List[ContextEnhancement], - team_knowledge: List[ContextEnhancement], - similar_tasks: List[ContextEnhancement], - ) -> List[str]: - """Generate actionable recommendations based on knowledge""" - - recommendations = [] - - # Recommendations from high-value organizational knowledge - high_value_org = [item for item in org_knowledge if item.confidence_score > 0.7] - for item in high_value_org[:3]: - if item.category == "best_practice": - recommendations.append(f"Apply best practice: {item.title}") - elif item.category == "development": - recommendations.append(f"Consider development approach: {item.title}") - - # Recommendations from effective team knowledge - effective_team = [ - item - for item in team_knowledge - if item.usage_stats.get("adoption_rate", 0) > 0.5 - ] - for item in effective_team[:2]: - recommendations.append(f"Team recommendation: {item.title}") - - # Recommendations from similar successful tasks - for task in similar_tasks: - if task.relevance_score > 0.6: - recommendations.append( - f"Based on similar task: Consider approach used in '{task.title}'" - ) - - return recommendations[:8] # Limit recommendations - - def _create_context_summary( - self, - task_data: Dict[str, Any], - org_knowledge: List[ContextEnhancement], - team_knowledge: List[ContextEnhancement], - success_patterns: List[str], - ) -> str: - """Create a summary of the enhanced context""" - - summary_parts = [] - - summary_parts.append(f"Enhanced context for: {task_data.get('title', 'Task')}") - - if org_knowledge: - summary_parts.append( - f"• {len(org_knowledge)} organizational knowledge items available" - ) - - if team_knowledge: - summary_parts.append( - f"• {len(team_knowledge)} team-specific insights included" - ) - - if success_patterns: - summary_parts.append( - f"• {len(success_patterns)} success patterns identified" - ) - summary_parts.append(f"Key patterns: {', '.join(success_patterns[:3])}") - - return "\n".join(summary_parts) - - def _extract_task_insights(self, task_data: Dict, task_result: Dict) -> str: - """Extract insights from a completed task""" - - insights = [] - - # Extract approach information - if task_result.get("iterations"): - insights.append(f"Completed in {task_result['iterations']} iterations") - - if task_result.get("pull_request_url"): - insights.append("Successfully created pull request") - - # Extract process information - if task_data.get("description"): - insights.append(f"Approach: {task_data['description'][:100]}...") - - return "\n".join(insights) - - async def _get_iteration_guidance( - self, organization_id: str, iteration_count: int - ) -> List[str]: - """Get guidance for high iteration count situations""" - - guidance = [] - - if iteration_count > 5: - # Search for guidance on complex tasks - complex_task_knowledge = await self.org_rag_manager.search_knowledge( - organization_id=organization_id, - query="complex task multiple iterations debugging", - limit=3, - min_similarity=0.3, - ) - - for result in complex_task_knowledge: - if "process" in result.knowledge.knowledge_category.value: - guidance.append(f"Process guidance: {result.knowledge.title}") - - return guidance - - async def _get_contextual_suggestions( - self, organization_id: str, team_id: str, current_context: Dict[str, Any] - ) -> List[str]: - """Get suggestions based on current execution context""" - - suggestions = [] - - # Context-specific suggestions based on current state - if current_context.get("error_count", 0) > 2: - suggestions.append( - "Consider reviewing error patterns in organizational knowledge" - ) - - if current_context.get("execution_time_minutes", 0) > 60: - suggestions.append("Look for optimization guidance from team knowledge") - - return suggestions - - async def _track_enhancement_usage(self, enhanced_context: EnhancedContext): - """Track usage of enhancement for analytics""" - - try: - async with self.pool.acquire() as conn: - # This would store enhancement usage data for analytics - # Placeholder for actual implementation - pass - except Exception as e: - logger.error(f"Error tracking enhancement usage: {e}") +""" +Context Enhancement Service for FuzeAgent + +This service enhances agent context with relevant organizational and team knowledge +before task execution. It provides intelligent knowledge injection based on +task type, agent capabilities, and historical success patterns. +""" + +import asyncio +import json +import logging +from dataclasses import dataclass +from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple + +import asyncpg + +from .organization_rag_manager import ( + ContentType, + KnowledgeCategory, + KnowledgeSearchResult, + OrganizationRAGManager, +) +from .team_knowledge_manager import TeamKnowledgeManager, TeamKnowledgeSearchResult + +logger = logging.getLogger(__name__) + + +@dataclass +class ContextEnhancement: + """Represents an enhancement to agent context""" + + knowledge_id: str + title: str + content: str + source_type: str # 'organization', 'team', 'agent' + category: str + relevance_score: float + confidence_score: float + usage_stats: Dict[str, Any] + metadata: Dict[str, Any] + + +@dataclass +class EnhancedContext: + """Enhanced context for agent task execution""" + + task_id: str + agent_id: str + team_id: str + organization_id: str + base_context: Dict[str, Any] + organizational_knowledge: List[ContextEnhancement] + team_knowledge: List[ContextEnhancement] + similar_task_insights: List[ContextEnhancement] + success_patterns: List[str] + common_pitfalls: List[str] + recommended_approaches: List[str] + context_summary: str + enhancement_metadata: Dict[str, Any] + + +class ContextEnhancementService: + """ + Enhances agent context with relevant organizational knowledge + to improve task execution success rates. + """ + + def __init__( + self, + database_url: str, + org_rag_manager: OrganizationRAGManager, + team_knowledge_manager: TeamKnowledgeManager, + ): + self.database_url = database_url + self.org_rag_manager = org_rag_manager + self.team_knowledge_manager = team_knowledge_manager + self.pool: Optional[asyncpg.Pool] = None + + # Configuration + self.max_org_knowledge_items = 5 + self.max_team_knowledge_items = 8 + self.max_similar_tasks = 3 + self.min_relevance_threshold = 0.4 + self.context_freshness_days = 90 + + # Statistics + self.enhancements_created = 0 + self.average_enhancement_score = 0.0 + self.knowledge_usage_tracking = {} + + async def initialize(self): + """Initialize the context enhancement service""" + logger.info("Initializing ContextEnhancementService") + + try: + self.pool = await asyncpg.create_pool( + self.database_url, min_size=1, max_size=5, command_timeout=60 + ) + + logger.info("ContextEnhancementService initialized successfully") + + except Exception as e: + logger.error(f"Failed to initialize ContextEnhancementService: {e}") + raise + + async def close(self): + """Close database connections""" + if self.pool: + await self.pool.close() + logger.info("ContextEnhancementService closed") + + async def enhance_agent_context( + self, + agent_id: str, + task_data: Dict[str, Any], + base_context: Optional[Dict[str, Any]] = None, + ) -> EnhancedContext: + """Enhance agent context with relevant organizational knowledge""" + + try: + # Get agent and team information + agent_info = await self._get_agent_info(agent_id) + if not agent_info: + raise ValueError(f"Agent {agent_id} not found") + + # Build search queries based on task data + search_queries = self._build_search_queries(task_data, agent_info) + + # Gather knowledge from different sources + org_knowledge = await self._gather_organizational_knowledge( + agent_info["organization_id"], + search_queries, + agent_id, + agent_info["team_id"], + ) + + team_knowledge = await self._gather_team_knowledge( + agent_info["team_id"], search_queries + ) + + similar_tasks = await self._find_similar_task_insights( + agent_info["organization_id"], task_data, agent_id + ) + + # Extract patterns and recommendations + success_patterns = await self._extract_success_patterns( + org_knowledge + team_knowledge + similar_tasks + ) + + pitfalls = await self._extract_common_pitfalls( + agent_info["organization_id"], task_data + ) + + recommendations = await self._generate_recommendations( + task_data, org_knowledge, team_knowledge, similar_tasks + ) + + # Create context summary + context_summary = self._create_context_summary( + task_data, org_knowledge, team_knowledge, success_patterns + ) + + # Build enhanced context + enhanced_context = EnhancedContext( + task_id=task_data.get("task_id", ""), + agent_id=agent_id, + team_id=agent_info["team_id"], + organization_id=agent_info["organization_id"], + base_context=base_context or {}, + organizational_knowledge=org_knowledge, + team_knowledge=team_knowledge, + similar_task_insights=similar_tasks, + success_patterns=success_patterns, + common_pitfalls=pitfalls, + recommended_approaches=recommendations, + context_summary=context_summary, + enhancement_metadata={ + "enhancement_timestamp": datetime.now().isoformat(), + "search_queries_used": search_queries, + "knowledge_sources_count": { + "organizational": len(org_knowledge), + "team": len(team_knowledge), + "similar_tasks": len(similar_tasks), + }, + "total_relevance_score": sum( + item.relevance_score + for item in org_knowledge + team_knowledge + similar_tasks + ), + "enhancement_version": "1.0", + }, + ) + + # Track enhancement usage + await self._track_enhancement_usage(enhanced_context) + + self.enhancements_created += 1 + + logger.info( + f"Enhanced context for agent {agent_id}: " + f"{len(org_knowledge)} org + {len(team_knowledge)} team + " + f"{len(similar_tasks)} similar task insights" + ) + + return enhanced_context + + except Exception as e: + logger.error(f"Error enhancing context for agent {agent_id}: {e}") + # Return minimal enhanced context on error + return EnhancedContext( + task_id=task_data.get("task_id", ""), + agent_id=agent_id, + team_id="", + organization_id="", + base_context=base_context or {}, + organizational_knowledge=[], + team_knowledge=[], + similar_task_insights=[], + success_patterns=[], + common_pitfalls=[], + recommended_approaches=[], + context_summary="Context enhancement failed - using minimal context", + enhancement_metadata={"error": str(e)}, + ) + + async def get_contextual_guidance( + self, + agent_id: str, + current_task_context: Dict[str, Any], + current_iteration: int = 1, + ) -> Dict[str, Any]: + """Get contextual guidance during task execution""" + + try: + agent_info = await self._get_agent_info(agent_id) + if not agent_info: + return {"guidance": [], "suggestions": []} + + # Build guidance based on current context + guidance_items = [] + + # Get iteration-specific guidance + if current_iteration > 3: + guidance_items.extend( + await self._get_iteration_guidance( + agent_info["organization_id"], current_iteration + ) + ) + + # Get context-specific suggestions + suggestions = await self._get_contextual_suggestions( + agent_info["organization_id"], + agent_info["team_id"], + current_task_context, + ) + + return { + "guidance": guidance_items, + "suggestions": suggestions, + "iteration": current_iteration, + "generated_at": datetime.now().isoformat(), + } + + except Exception as e: + logger.error(f"Error getting contextual guidance: {e}") + return {"guidance": [], "suggestions": []} + + async def update_knowledge_effectiveness( + self, + knowledge_id: str, + knowledge_source: str, + task_success: bool, + agent_feedback: Optional[Dict[str, Any]] = None, + ): + """Update knowledge effectiveness based on usage outcomes""" + + try: + if knowledge_source == "organization": + await self.org_rag_manager.update_knowledge_quality( + knowledge_id=knowledge_id, + success_correlation=1.0 if task_success else -0.2, + feedback_metadata=agent_feedback, + ) + elif knowledge_source == "team": + # Update team knowledge effectiveness + async with self.pool.acquire() as conn: + agent_id = ( + agent_feedback.get("agent_id") if agent_feedback else None + ) + if agent_id: + await self.team_knowledge_manager.update_knowledge_effectiveness( + team_knowledge_id=knowledge_id, + agent_id=agent_id, + task_success=task_success, + feedback_score=agent_feedback.get("usefulness_score"), + usage_context=agent_feedback, + ) + + # Track in local usage statistics + if knowledge_id not in self.knowledge_usage_tracking: + self.knowledge_usage_tracking[knowledge_id] = { + "usage_count": 0, + "success_count": 0, + "effectiveness_score": 0.0, + } + + stats = self.knowledge_usage_tracking[knowledge_id] + stats["usage_count"] += 1 + if task_success: + stats["success_count"] += 1 + stats["effectiveness_score"] = stats["success_count"] / stats["usage_count"] + + except Exception as e: + logger.error(f"Error updating knowledge effectiveness: {e}") + + async def get_enhancement_statistics( + self, + organization_id: Optional[str] = None, + team_id: Optional[str] = None, + days_back: int = 30, + ) -> Dict[str, Any]: + """Get context enhancement statistics""" + + try: + async with self.pool.acquire() as conn: + # Basic enhancement statistics + where_conditions = [ + "created_at >= NOW() - INTERVAL '%s days'" % days_back + ] + params = [] + + if organization_id: + where_conditions.append("organization_id = $1") + params.append(organization_id) + + if team_id: + where_conditions.append( + "team_id = $2" if organization_id else "team_id = $1" + ) + params.append(team_id) + + # This is a placeholder - in practice you'd have a table to track enhancements + stats = { + "total_enhancements": self.enhancements_created, + "average_knowledge_items_per_enhancement": { + "organizational": 3.2, + "team": 4.1, + "similar_tasks": 1.8, + }, + "knowledge_effectiveness": dict(self.knowledge_usage_tracking), + "generated_at": datetime.now().isoformat(), + } + + return stats + + except Exception as e: + logger.error(f"Error getting enhancement statistics: {e}") + return {} + + async def _get_agent_info(self, agent_id: str) -> Optional[Dict[str, Any]]: + """Get agent information including team and organization""" + + async with self.pool.acquire() as conn: + agent_info = await conn.fetchrow( + """ + SELECT a.id, a.name, a.type, a.config, a.team_id, + t.organization_id, t.name as team_name, + o.name as organization_name + FROM agents a + JOIN teams t ON a.team_id = t.id + JOIN organizations o ON t.organization_id = o.id + WHERE a.id = $1 + """, + agent_id, + ) + + return dict(agent_info) if agent_info else None + + def _build_search_queries( + self, task_data: Dict[str, Any], agent_info: Dict[str, Any] + ) -> List[str]: + """Build search queries based on task data and agent information""" + + queries = [] + + # Primary query from task description + if task_data.get("description"): + queries.append(task_data["description"][:200]) + + # Query from task title + if task_data.get("title"): + queries.append(task_data["title"]) + + # Technology-specific queries + if task_data.get("technologies"): + for tech in task_data["technologies"]: + queries.append(f"{tech} development best practices") + + # Task type specific query + if task_data.get("task_type"): + queries.append(f"{task_data['task_type']} implementation guide") + + # Agent type specific query + agent_type = agent_info.get("type", "") + if agent_type: + queries.append(f"{agent_type} workflow best practices") + + return queries[:5] # Limit to top 5 queries + + async def _gather_organizational_knowledge( + self, + organization_id: str, + search_queries: List[str], + agent_id: str, + team_id: str, + ) -> List[ContextEnhancement]: + """Gather relevant organizational knowledge""" + + org_knowledge = [] + + for query in search_queries: + search_results = await self.org_rag_manager.search_knowledge( + organization_id=organization_id, + query=query, + limit=self.max_org_knowledge_items // len(search_queries) + 1, + min_similarity=self.min_relevance_threshold, + requester_agent_id=agent_id, + requester_team_id=team_id, + ) + + for result in search_results: + if result.combined_score >= self.min_relevance_threshold: + enhancement = ContextEnhancement( + knowledge_id=result.knowledge.id, + title=result.knowledge.title, + content=( + result.knowledge.content[:1000] + "..." + if len(result.knowledge.content) > 1000 + else result.knowledge.content + ), + source_type="organization", + category=result.knowledge.knowledge_category.value, + relevance_score=result.combined_score, + confidence_score=result.knowledge.quality_score, + usage_stats={ + "usage_count": result.knowledge.usage_count, + "success_correlation": result.knowledge.success_correlation, + }, + metadata=result.knowledge.metadata, + ) + org_knowledge.append(enhancement) + + # Sort by relevance and remove duplicates + seen_ids = set() + unique_knowledge = [] + for item in sorted( + org_knowledge, key=lambda x: x.relevance_score, reverse=True + ): + if item.knowledge_id not in seen_ids: + unique_knowledge.append(item) + seen_ids.add(item.knowledge_id) + + return unique_knowledge[: self.max_org_knowledge_items] + + async def _gather_team_knowledge( + self, team_id: str, search_queries: List[str] + ) -> List[ContextEnhancement]: + """Gather relevant team knowledge""" + + team_knowledge = [] + + for query in search_queries: + search_results = await self.team_knowledge_manager.search_team_knowledge( + team_id=team_id, + query=query, + limit=self.max_team_knowledge_items // len(search_queries) + 1, + min_similarity=self.min_relevance_threshold, + ) + + for result in search_results: + if result.combined_score >= self.min_relevance_threshold: + enhancement = ContextEnhancement( + knowledge_id=result.team_knowledge.id, + title=result.team_knowledge.title, + content=( + result.team_knowledge.content[:1000] + "..." + if len(result.team_knowledge.content) > 1000 + else result.team_knowledge.content + ), + source_type="team", + category=result.team_knowledge.knowledge_category.value, + relevance_score=result.combined_score, + confidence_score=result.team_knowledge.effectiveness_score, + usage_stats={ + "adoption_rate": result.team_knowledge.agent_adoption_rate, + "effectiveness": result.team_knowledge.effectiveness_score, + }, + metadata=result.team_knowledge.metadata, + ) + team_knowledge.append(enhancement) + + # Sort and deduplicate + seen_ids = set() + unique_knowledge = [] + for item in sorted( + team_knowledge, key=lambda x: x.relevance_score, reverse=True + ): + if item.knowledge_id not in seen_ids: + unique_knowledge.append(item) + seen_ids.add(item.knowledge_id) + + return unique_knowledge[: self.max_team_knowledge_items] + + async def _find_similar_task_insights( + self, organization_id: str, task_data: Dict[str, Any], agent_id: str + ) -> List[ContextEnhancement]: + """Find insights from similar completed tasks""" + + similar_tasks = [] + + try: + async with self.pool.acquire() as conn: + # Find similar tasks based on description similarity and success + similar_task_data = await conn.fetch( + """ + SELECT t.id, t.title, t.description, t.result, t.completed_at, + a.name as agent_name, a.type as agent_type, + similarity(t.description, $2) as similarity_score + FROM tasks t + JOIN agents a ON t.agent_id = a.id + JOIN teams te ON a.team_id = te.id + WHERE te.organization_id = $1 + AND t.status = 'completed' + AND t.result->>'status' = 'completed' + AND t.completed_at >= NOW() - INTERVAL '90 days' + AND similarity(t.description, $2) > 0.3 + ORDER BY similarity_score DESC, t.completed_at DESC + LIMIT $3 + """, + organization_id, + task_data.get("description", ""), + self.max_similar_tasks, + ) + + for task in similar_task_data: + # Extract insights from the task result + task_result = ( + task["result"] if isinstance(task["result"], dict) else {} + ) + + insights_content = self._extract_task_insights( + dict(task), task_result + ) + + if insights_content: + enhancement = ContextEnhancement( + knowledge_id=str(task["id"]), + title=f"Similar Task: {task['title'][:50]}...", + content=insights_content, + source_type="similar_task", + category="process", + relevance_score=float(task["similarity_score"]), + confidence_score=0.8, # High confidence for successful completed tasks + usage_stats={ + "agent_type": task["agent_type"], + "completion_date": task["completed_at"].isoformat(), + }, + metadata={ + "source_task_id": str(task["id"]), + "source_agent": task["agent_name"], + "similarity_score": float(task["similarity_score"]), + }, + ) + similar_tasks.append(enhancement) + + except Exception as e: + logger.error(f"Error finding similar task insights: {e}") + + return similar_tasks + + async def _extract_success_patterns( + self, all_knowledge: List[ContextEnhancement] + ) -> List[str]: + """Extract success patterns from knowledge items""" + + patterns = [] + + for item in all_knowledge: + # Look for success indicators in metadata + if "success_indicators" in item.metadata: + patterns.extend(item.metadata["success_indicators"]) + + # Extract patterns from high-confidence, high-usage items + if item.confidence_score > 0.7 and item.relevance_score > 0.6: + if "optimization" in item.title.lower(): + patterns.append("Focus on optimization early") + if "test" in item.title.lower(): + patterns.append("Comprehensive testing leads to success") + if "pattern" in item.title.lower(): + patterns.append("Follow established patterns") + + return list(set(patterns)) # Remove duplicates + + async def _extract_common_pitfalls( + self, organization_id: str, task_data: Dict[str, Any] + ) -> List[str]: + """Extract common pitfalls for this type of task""" + + pitfalls = [] + + try: + # Search for error patterns and failure knowledge + error_knowledge = await self.org_rag_manager.search_knowledge( + organization_id=organization_id, + query=f"error pattern {task_data.get('task_type', '')}", + categories=[KnowledgeCategory.TROUBLESHOOTING], + limit=5, + min_similarity=0.3, + ) + + for result in error_knowledge: + if "failure_patterns" in result.knowledge.metadata: + pitfalls.extend(result.knowledge.metadata["failure_patterns"]) + + # Extract pitfalls from error pattern content + content_lower = result.knowledge.content.lower() + if ( + "avoid" in content_lower + or "pitfall" in content_lower + or "common mistake" in content_lower + ): + pitfalls.append(result.knowledge.title) + + except Exception as e: + logger.error(f"Error extracting pitfalls: {e}") + + return list(set(pitfalls))[:5] # Top 5 pitfalls + + async def _generate_recommendations( + self, + task_data: Dict[str, Any], + org_knowledge: List[ContextEnhancement], + team_knowledge: List[ContextEnhancement], + similar_tasks: List[ContextEnhancement], + ) -> List[str]: + """Generate actionable recommendations based on knowledge""" + + recommendations = [] + + # Recommendations from high-value organizational knowledge + high_value_org = [item for item in org_knowledge if item.confidence_score > 0.7] + for item in high_value_org[:3]: + if item.category == "best_practice": + recommendations.append(f"Apply best practice: {item.title}") + elif item.category == "development": + recommendations.append(f"Consider development approach: {item.title}") + + # Recommendations from effective team knowledge + effective_team = [ + item + for item in team_knowledge + if item.usage_stats.get("adoption_rate", 0) > 0.5 + ] + for item in effective_team[:2]: + recommendations.append(f"Team recommendation: {item.title}") + + # Recommendations from similar successful tasks + for task in similar_tasks: + if task.relevance_score > 0.6: + recommendations.append( + f"Based on similar task: Consider approach used in '{task.title}'" + ) + + return recommendations[:8] # Limit recommendations + + def _create_context_summary( + self, + task_data: Dict[str, Any], + org_knowledge: List[ContextEnhancement], + team_knowledge: List[ContextEnhancement], + success_patterns: List[str], + ) -> str: + """Create a summary of the enhanced context""" + + summary_parts = [] + + summary_parts.append(f"Enhanced context for: {task_data.get('title', 'Task')}") + + if org_knowledge: + summary_parts.append( + f"• {len(org_knowledge)} organizational knowledge items available" + ) + + if team_knowledge: + summary_parts.append( + f"• {len(team_knowledge)} team-specific insights included" + ) + + if success_patterns: + summary_parts.append( + f"• {len(success_patterns)} success patterns identified" + ) + summary_parts.append(f"Key patterns: {', '.join(success_patterns[:3])}") + + return "\n".join(summary_parts) + + def _extract_task_insights(self, task_data: Dict, task_result: Dict) -> str: + """Extract insights from a completed task""" + + insights = [] + + # Extract approach information + if task_result.get("iterations"): + insights.append(f"Completed in {task_result['iterations']} iterations") + + if task_result.get("pull_request_url"): + insights.append("Successfully created pull request") + + # Extract process information + if task_data.get("description"): + insights.append(f"Approach: {task_data['description'][:100]}...") + + return "\n".join(insights) + + async def _get_iteration_guidance( + self, organization_id: str, iteration_count: int + ) -> List[str]: + """Get guidance for high iteration count situations""" + + guidance = [] + + if iteration_count > 5: + # Search for guidance on complex tasks + complex_task_knowledge = await self.org_rag_manager.search_knowledge( + organization_id=organization_id, + query="complex task multiple iterations debugging", + limit=3, + min_similarity=0.3, + ) + + for result in complex_task_knowledge: + if "process" in result.knowledge.knowledge_category.value: + guidance.append(f"Process guidance: {result.knowledge.title}") + + return guidance + + async def _get_contextual_suggestions( + self, organization_id: str, team_id: str, current_context: Dict[str, Any] + ) -> List[str]: + """Get suggestions based on current execution context""" + + suggestions = [] + + # Context-specific suggestions based on current state + if current_context.get("error_count", 0) > 2: + suggestions.append( + "Consider reviewing error patterns in organizational knowledge" + ) + + if current_context.get("execution_time_minutes", 0) > 60: + suggestions.append("Look for optimization guidance from team knowledge") + + return suggestions + + async def _track_enhancement_usage(self, enhanced_context: EnhancedContext): + """Track usage of enhancement for analytics""" + + try: + async with self.pool.acquire() as conn: + # This would store enhancement usage data for analytics + # Placeholder for actual implementation + pass + except Exception as e: + logger.error(f"Error tracking enhancement usage: {e}") diff --git a/services/orchestrator/context_service.py b/services/orchestrator/context_service.py index 7dd7bab..e7092b4 100644 --- a/services/orchestrator/context_service.py +++ b/services/orchestrator/context_service.py @@ -1,152 +1,152 @@ -import asyncio -import json -from typing import Dict, List, Optional - -import numpy as np -from sentence_transformers import SentenceTransformer - -from .database import get_db_connection - - -class ContextService: - def __init__(self): - # Load sentence transformer for embeddings - self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") - - def generate_embedding(self, text: str) -> List[float]: - """Generate embedding for text""" - embedding = self.embedding_model.encode(text) - return embedding.tolist() - - async def store_interaction( - self, agent_id: str, content: str, metadata: Dict = None - ) -> str: - """Store agent interaction with embedding""" - embedding = self.generate_embedding(content) - - async with get_db_connection() as conn: - interaction_id = await conn.fetchval( - """ - INSERT INTO interactions (agent_id, content, embedding, metadata) - VALUES ($1, $2, $3, $4) - RETURNING id - """, - agent_id, - content, - embedding, - metadata or {}, - ) - return str(interaction_id) - - async def get_similar_interactions( - self, - query: str, - agent_id: str = None, - limit: int = 5, - similarity_threshold: float = 0.7, - ) -> List[Dict]: - """Find similar interactions using vector similarity""" - query_embedding = self.generate_embedding(query) - - async with get_db_connection() as conn: - if agent_id: - rows = await conn.fetch( - """ - SELECT id, agent_id, content, metadata, created_at, - 1 - (embedding <=> $1) as similarity - FROM interactions - WHERE agent_id = $2 AND 1 - (embedding <=> $1) > $3 - ORDER BY similarity DESC - LIMIT $4 - """, - query_embedding, - agent_id, - similarity_threshold, - limit, - ) - else: - rows = await conn.fetch( - """ - SELECT id, agent_id, content, metadata, created_at, - 1 - (embedding <=> $1) as similarity - FROM interactions - WHERE 1 - (embedding <=> $1) > $2 - ORDER BY similarity DESC - LIMIT $3 - """, - query_embedding, - similarity_threshold, - limit, - ) - - return [dict(row) for row in rows] - - async def get_context(self, query: str, agent_id: str = None) -> Dict: - """Get relevant context for a query""" - similar_interactions = await self.get_similar_interactions(query, agent_id) - - # Get recent interactions from same agent - recent_interactions = [] - if agent_id: - async with get_db_connection() as conn: - rows = await conn.fetch( - """ - SELECT content, metadata, created_at - FROM interactions - WHERE agent_id = $1 - ORDER BY created_at DESC - LIMIT 10 - """, - agent_id, - ) - recent_interactions = [dict(row) for row in rows] - - return { - "similar_interactions": similar_interactions, - "recent_interactions": recent_interactions, - "relevant_code": self._extract_code_snippets(similar_interactions), - "similar_features": self._extract_similar_features(similar_interactions), - } - - def _extract_code_snippets(self, interactions: List[Dict]) -> str: - """Extract code snippets from interactions""" - code_snippets = [] - for interaction in interactions: - content = interaction.get("content", "") - # Simple extraction - look for code blocks - if "```" in content: - parts = content.split("```") - for i in range(1, len(parts), 2): - code_snippets.append(parts[i].strip()) - - return "\n\n".join(code_snippets[:3]) # Return top 3 snippets - - def _extract_similar_features(self, interactions: List[Dict]) -> str: - """Extract similar feature descriptions""" - features = [] - for interaction in interactions: - metadata = interaction.get("metadata", {}) - if "task_type" in metadata and metadata["task_type"] == "implement_feature": - features.append(interaction.get("content", "")) - - return "\n\n".join(features[:2]) # Return top 2 similar features - - async def get_agent_memory(self, agent_id: str, limit: int = 50) -> List[Dict]: - """Get agent's memory/interaction history""" - async with get_db_connection() as conn: - rows = await conn.fetch( - """ - SELECT content, metadata, created_at - FROM interactions - WHERE agent_id = $1 - ORDER BY created_at DESC - LIMIT $2 - """, - agent_id, - limit, - ) - return [dict(row) for row in rows] - - async def search_knowledge(self, query: str, limit: int = 10) -> List[Dict]: - """Search across all agent knowledge""" - return await self.get_similar_interactions(query, limit=limit) +import asyncio +import json +from typing import Dict, List, Optional + +import numpy as np +from sentence_transformers import SentenceTransformer + +from .database import get_db_connection + + +class ContextService: + def __init__(self): + # Load sentence transformer for embeddings + self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") + + def generate_embedding(self, text: str) -> List[float]: + """Generate embedding for text""" + embedding = self.embedding_model.encode(text) + return embedding.tolist() + + async def store_interaction( + self, agent_id: str, content: str, metadata: Dict = None + ) -> str: + """Store agent interaction with embedding""" + embedding = self.generate_embedding(content) + + async with get_db_connection() as conn: + interaction_id = await conn.fetchval( + """ + INSERT INTO interactions (agent_id, content, embedding, metadata) + VALUES ($1, $2, $3, $4) + RETURNING id + """, + agent_id, + content, + embedding, + metadata or {}, + ) + return str(interaction_id) + + async def get_similar_interactions( + self, + query: str, + agent_id: str = None, + limit: int = 5, + similarity_threshold: float = 0.7, + ) -> List[Dict]: + """Find similar interactions using vector similarity""" + query_embedding = self.generate_embedding(query) + + async with get_db_connection() as conn: + if agent_id: + rows = await conn.fetch( + """ + SELECT id, agent_id, content, metadata, created_at, + 1 - (embedding <=> $1) as similarity + FROM interactions + WHERE agent_id = $2 AND 1 - (embedding <=> $1) > $3 + ORDER BY similarity DESC + LIMIT $4 + """, + query_embedding, + agent_id, + similarity_threshold, + limit, + ) + else: + rows = await conn.fetch( + """ + SELECT id, agent_id, content, metadata, created_at, + 1 - (embedding <=> $1) as similarity + FROM interactions + WHERE 1 - (embedding <=> $1) > $2 + ORDER BY similarity DESC + LIMIT $3 + """, + query_embedding, + similarity_threshold, + limit, + ) + + return [dict(row) for row in rows] + + async def get_context(self, query: str, agent_id: str = None) -> Dict: + """Get relevant context for a query""" + similar_interactions = await self.get_similar_interactions(query, agent_id) + + # Get recent interactions from same agent + recent_interactions = [] + if agent_id: + async with get_db_connection() as conn: + rows = await conn.fetch( + """ + SELECT content, metadata, created_at + FROM interactions + WHERE agent_id = $1 + ORDER BY created_at DESC + LIMIT 10 + """, + agent_id, + ) + recent_interactions = [dict(row) for row in rows] + + return { + "similar_interactions": similar_interactions, + "recent_interactions": recent_interactions, + "relevant_code": self._extract_code_snippets(similar_interactions), + "similar_features": self._extract_similar_features(similar_interactions), + } + + def _extract_code_snippets(self, interactions: List[Dict]) -> str: + """Extract code snippets from interactions""" + code_snippets = [] + for interaction in interactions: + content = interaction.get("content", "") + # Simple extraction - look for code blocks + if "```" in content: + parts = content.split("```") + for i in range(1, len(parts), 2): + code_snippets.append(parts[i].strip()) + + return "\n\n".join(code_snippets[:3]) # Return top 3 snippets + + def _extract_similar_features(self, interactions: List[Dict]) -> str: + """Extract similar feature descriptions""" + features = [] + for interaction in interactions: + metadata = interaction.get("metadata", {}) + if "task_type" in metadata and metadata["task_type"] == "implement_feature": + features.append(interaction.get("content", "")) + + return "\n\n".join(features[:2]) # Return top 2 similar features + + async def get_agent_memory(self, agent_id: str, limit: int = 50) -> List[Dict]: + """Get agent's memory/interaction history""" + async with get_db_connection() as conn: + rows = await conn.fetch( + """ + SELECT content, metadata, created_at + FROM interactions + WHERE agent_id = $1 + ORDER BY created_at DESC + LIMIT $2 + """, + agent_id, + limit, + ) + return [dict(row) for row in rows] + + async def search_knowledge(self, query: str, limit: int = 10) -> List[Dict]: + """Search across all agent knowledge""" + return await self.get_similar_interactions(query, limit=limit) diff --git a/services/orchestrator/conversation_manager.py b/services/orchestrator/conversation_manager.py index 507aa26..e792c2c 100644 --- a/services/orchestrator/conversation_manager.py +++ b/services/orchestrator/conversation_manager.py @@ -1,602 +1,602 @@ -""" -Conversation Manager for FuzeAgent Claude Code Integration - -Manages and stores complete conversations between agents and Claude Code, -providing comprehensive audit trails, debugging capabilities, and learning data. -""" - -import asyncio -import json -import logging -import time -import uuid -from dataclasses import dataclass -from datetime import datetime -from enum import Enum -from typing import Any, Dict, List, Optional - -# Importable both as `services.orchestrator.conversation_manager` (relative) and -# flat with services/orchestrator on sys.path (as the tests import it). -try: - from .database import get_db_connection -except ImportError: # pragma: no cover - flat import (no parent package) - from database import get_db_connection - -logger = logging.getLogger(__name__) - - -class MessageType(str, Enum): - USER_PROMPT = "user_prompt" - CLAUDE_RESPONSE = "claude_response" - SYSTEM_MESSAGE = "system_message" - ERROR_MESSAGE = "error_message" - CODE_EXECUTION = "code_execution" - TEST_RESULT = "test_result" - - -class InteractionType(str, Enum): - QUESTION = "question" - CLARIFICATION = "clarification" - APPROVAL_REQUEST = "approval_request" - ERROR_REPORT = "error_report" - PROGRESS_UPDATE = "progress_update" - - -@dataclass -class ConversationMessage: - """Represents a single message in a Claude Code conversation""" - - task_id: str - iteration_number: int - message_type: MessageType - content: str - token_count: Optional[int] = None - model_used: Optional[str] = None - temperature: Optional[float] = None - response_time_ms: Optional[int] = None - metadata: Optional[Dict[str, Any]] = None - - -@dataclass -class ConversationSession: - """Represents a complete conversation session""" - - agent_id: str - task_id: str - sandbox_id: str - session_started_at: datetime - session_ended_at: Optional[datetime] = None - total_messages: int = 0 - total_tokens: int = 0 - status: str = "active" - metadata: Optional[Dict[str, Any]] = None - - -class ConversationManager: - """ - Manages Claude Code conversations and provides comprehensive tracking. - - Features: - - Full conversation storage and retrieval - - Token usage tracking and cost analysis - - Human interaction management - - Code generation tracking - - Performance metrics collection - """ - - def __init__(self): - self.active_sessions: Dict[str, ConversationSession] = {} - - async def start_conversation_session( - self, - agent_id: str, - task_id: str, - sandbox_id: str, - metadata: Optional[Dict[str, Any]] = None, - ) -> str: - """Start a new conversation session for an agent""" - - session = ConversationSession( - agent_id=agent_id, - task_id=task_id, - sandbox_id=sandbox_id, - session_started_at=datetime.now(), - metadata=metadata or {}, - ) - - # Store session in database - session_id = await self._store_conversation_session(session) - session.metadata = session.metadata or {} - session.metadata["session_id"] = session_id - - # Track active session - self.active_sessions[session_id] = session - - logger.info( - f"Started conversation session {session_id} for agent {agent_id}, task {task_id}" - ) - return session_id - - async def end_conversation_session(self, session_id: str) -> bool: - """End a conversation session""" - - session = self.active_sessions.get(session_id) - if not session: - logger.warning(f"Session {session_id} not found in active sessions") - return False - - session.session_ended_at = datetime.now() - session.status = "completed" - - # Update database - await self._update_conversation_session(session_id, session) - - # Remove from active sessions - del self.active_sessions[session_id] - - logger.info(f"Ended conversation session {session_id}") - return True - - async def store_message( - self, - session_id: str, - message: ConversationMessage, - start_time: Optional[float] = None, - ) -> str: - """Store a conversation message""" - - # Calculate response time if start_time provided - if start_time and message.message_type == MessageType.CLAUDE_RESPONSE: - message.response_time_ms = int((time.time() - start_time) * 1000) - - # Store message in database - message_id = await self._store_claude_conversation(message) - - # Update session statistics - session = self.active_sessions.get(session_id) - if session: - session.total_messages += 1 - if message.token_count: - session.total_tokens += message.token_count - - logger.debug(f"Stored message {message_id} for session {session_id}") - return message_id - - async def store_user_prompt( - self, - session_id: str, - task_id: str, - iteration_number: int, - prompt: str, - model: str = "claude-3-5-sonnet-20241022", - temperature: float = 0.3, - metadata: Optional[Dict[str, Any]] = None, - ) -> str: - """Store a user prompt to Claude""" - - message = ConversationMessage( - task_id=task_id, - iteration_number=iteration_number, - message_type=MessageType.USER_PROMPT, - content=prompt, - model_used=model, - temperature=temperature, - metadata=metadata or {}, - ) - - return await self.store_message(session_id, message) - - async def store_claude_response( - self, - session_id: str, - task_id: str, - iteration_number: int, - response: str, - token_count: Optional[int] = None, - model: str = "claude-3-5-sonnet-20241022", - start_time: Optional[float] = None, - metadata: Optional[Dict[str, Any]] = None, - ) -> str: - """Store Claude's response""" - - message = ConversationMessage( - task_id=task_id, - iteration_number=iteration_number, - message_type=MessageType.CLAUDE_RESPONSE, - content=response, - token_count=token_count, - model_used=model, - metadata=metadata or {}, - ) - - return await self.store_message(session_id, message, start_time) - - async def store_code_generation( - self, - task_id: str, - iteration_number: int, - file_path: str, - file_type: str, - language: str, - content: str, - commit_hash: Optional[str] = None, - test_results: Optional[Dict[str, Any]] = None, - quality_metrics: Optional[Dict[str, Any]] = None, - ) -> str: - """Store generated code with metadata""" - - async with get_db_connection() as conn: - code_id = await conn.fetchval( - """ - INSERT INTO code_generations ( - task_id, iteration_number, file_path, file_type, language, - content, commit_hash, test_results, quality_metrics - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) - RETURNING id - """, - task_id, - iteration_number, - file_path, - file_type, - language, - content, - commit_hash, - json.dumps(test_results) if test_results else None, - json.dumps(quality_metrics) if quality_metrics else None, - ) - - logger.info(f"Stored code generation {code_id} for task {task_id}") - return str(code_id) - - async def store_human_interaction( - self, - task_id: str, - iteration_number: int, - interaction_type: InteractionType, - agent_message: str, - human_response: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, - ) -> str: - """Store human-in-the-loop interaction""" - - async with get_db_connection() as conn: - interaction_id = await conn.fetchval( - """ - INSERT INTO human_interactions ( - task_id, iteration_number, interaction_type, - agent_message, human_response, metadata - ) VALUES ($1, $2, $3, $4, $5, $6) - RETURNING id - """, - task_id, - iteration_number, - interaction_type.value, - agent_message, - human_response, - json.dumps(metadata) if metadata else {}, - ) - - logger.info(f"Stored human interaction {interaction_id} for task {task_id}") - return str(interaction_id) - - async def update_human_response( - self, interaction_id: str, human_response: str - ) -> bool: - """Update human response to an interaction""" - - async with get_db_connection() as conn: - result = await conn.execute( - """ - UPDATE human_interactions - SET human_response = $1, - responded_at = CURRENT_TIMESTAMP, - response_time_seconds = EXTRACT(EPOCH FROM (CURRENT_TIMESTAMP - asked_at)) - WHERE id = $2 - """, - human_response, - interaction_id, - ) - - success = result != "UPDATE 0" - if success: - logger.info(f"Updated human response for interaction {interaction_id}") - return success - - async def store_performance_metric( - self, - agent_id: str, - task_id: str, - metric_type: str, - metric_value: float, - metric_unit: Optional[str] = None, - context: Optional[Dict[str, Any]] = None, - ) -> str: - """Store agent performance metric""" - - async with get_db_connection() as conn: - metric_id = await conn.fetchval( - """ - INSERT INTO agent_performance_metrics ( - agent_id, task_id, metric_type, metric_value, - metric_unit, context - ) VALUES ($1, $2, $3, $4, $5, $6) - RETURNING id - """, - agent_id, - task_id, - metric_type, - metric_value, - metric_unit, - json.dumps(context) if context else {}, - ) - - logger.debug( - f"Stored performance metric {metric_id}: {metric_type}={metric_value}" - ) - return str(metric_id) - - async def get_conversation_history( - self, - task_id: str, - iteration_number: Optional[int] = None, - message_types: Optional[List[MessageType]] = None, - limit: Optional[int] = None, - ) -> List[Dict[str, Any]]: - """Get conversation history for a task""" - - conditions = ["task_id = $1"] - params = [task_id] - param_count = 1 - - if iteration_number is not None: - param_count += 1 - conditions.append(f"iteration_number = ${param_count}") - params.append(iteration_number) - - if message_types: - param_count += 1 - conditions.append(f"message_type = ANY(${param_count})") - params.append([mt.value for mt in message_types]) - - where_clause = " AND ".join(conditions) - limit_clause = "" - if limit: - param_count += 1 - limit_clause = f"LIMIT ${param_count}" - params.append(limit) - - async with get_db_connection() as conn: - rows = await conn.fetch( - f""" - SELECT * FROM claude_conversations - WHERE {where_clause} - ORDER BY created_at ASC - {limit_clause} - """, # nosec B608 -- where/limit clauses are fixed fragments with $N placeholders; all values bound as query params - *params, - ) - - return [dict(row) for row in rows] - - async def get_conversation_summary(self, task_id: str) -> Dict[str, Any]: - """Get conversation summary with statistics""" - - async with get_db_connection() as conn: - # Get message statistics - stats = await conn.fetchrow( - """ - SELECT - COUNT(*) as total_messages, - SUM(token_count) as total_tokens, - AVG(response_time_ms) as avg_response_time, - MAX(iteration_number) as max_iteration - FROM claude_conversations - WHERE task_id = $1 - """, - task_id, - ) - - # Get message type breakdown - type_breakdown = await conn.fetch( - """ - SELECT message_type, COUNT(*) as count - FROM claude_conversations - WHERE task_id = $1 - GROUP BY message_type - """, - task_id, - ) - - # Get human interactions - human_interactions = await conn.fetchrow( - """ - SELECT - COUNT(*) as total_interactions, - COUNT(human_response) as responded_interactions, - AVG(response_time_seconds) as avg_response_time - FROM human_interactions - WHERE task_id = $1 - """, - task_id, - ) - - return { - "task_id": task_id, - "total_messages": stats["total_messages"] or 0, - "total_tokens": stats["total_tokens"] or 0, - "avg_response_time_ms": float(stats["avg_response_time"] or 0), - "max_iteration": stats["max_iteration"] or 0, - "message_types": { - row["message_type"]: row["count"] for row in type_breakdown - }, - "human_interactions": { - "total": human_interactions["total_interactions"] or 0, - "responded": human_interactions["responded_interactions"] or 0, - "avg_response_time_seconds": float( - human_interactions["avg_response_time"] or 0 - ), - }, - } - - async def get_code_generations( - self, - task_id: str, - iteration_number: Optional[int] = None, - file_type: Optional[str] = None, - language: Optional[str] = None, - ) -> List[Dict[str, Any]]: - """Get code generations for a task""" - - conditions = ["task_id = $1"] - params = [task_id] - param_count = 1 - - if iteration_number is not None: - param_count += 1 - conditions.append(f"iteration_number = ${param_count}") - params.append(iteration_number) - - if file_type: - param_count += 1 - conditions.append(f"file_type = ${param_count}") - params.append(file_type) - - if language: - param_count += 1 - conditions.append(f"language = ${param_count}") - params.append(language) - - where_clause = " AND ".join(conditions) - - async with get_db_connection() as conn: - rows = await conn.fetch( - f""" - SELECT * FROM code_generations - WHERE {where_clause} - ORDER BY generated_at ASC - """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params - *params, - ) - - return [dict(row) for row in rows] - - async def get_agent_performance_metrics( - self, - agent_id: Optional[str] = None, - task_id: Optional[str] = None, - metric_types: Optional[List[str]] = None, - time_range_hours: Optional[int] = None, - ) -> List[Dict[str, Any]]: - """Get agent performance metrics""" - - conditions = [] - params = [] - param_count = 0 - - if agent_id: - param_count += 1 - conditions.append(f"agent_id = ${param_count}") - params.append(agent_id) - - if task_id: - param_count += 1 - conditions.append(f"task_id = ${param_count}") - params.append(task_id) - - if metric_types: - param_count += 1 - conditions.append(f"metric_type = ANY(${param_count})") - params.append(metric_types) - - if time_range_hours: - param_count += 1 - conditions.append( - f"measured_at >= NOW() - (INTERVAL '1 hour' * ${param_count})" - ) - params.append(time_range_hours) - - where_clause = "WHERE " + " AND ".join(conditions) if conditions else "" - - async with get_db_connection() as conn: - rows = await conn.fetch( - f""" - SELECT * FROM agent_performance_metrics - {where_clause} - ORDER BY measured_at DESC - """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values (incl. interval multiplier) bound as query params - *params, - ) - - return [dict(row) for row in rows] - - # Private methods - - async def _store_conversation_session(self, session: ConversationSession) -> str: - """Store conversation session in database""" - - async with get_db_connection() as conn: - session_id = await conn.fetchval( - """ - INSERT INTO agent_conversation_sessions ( - agent_id, task_id, sandbox_id, session_started_at, - total_messages, total_tokens, status, metadata - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) - RETURNING id - """, - session.agent_id, - session.task_id, - session.sandbox_id, - session.session_started_at, - session.total_messages, - session.total_tokens, - session.status, - json.dumps(session.metadata) if session.metadata else {}, - ) - - return str(session_id) - - async def _update_conversation_session( - self, session_id: str, session: ConversationSession - ): - """Update conversation session in database""" - - async with get_db_connection() as conn: - await conn.execute( - """ - UPDATE agent_conversation_sessions - SET session_ended_at = $1, total_messages = $2, - total_tokens = $3, status = $4, metadata = $5 - WHERE id = $6 - """, - session.session_ended_at, - session.total_messages, - session.total_tokens, - session.status, - json.dumps(session.metadata) if session.metadata else {}, - session_id, - ) - - async def _store_claude_conversation(self, message: ConversationMessage) -> str: - """Store Claude conversation message in database""" - - async with get_db_connection() as conn: - message_id = await conn.fetchval( - """ - INSERT INTO claude_conversations ( - task_id, iteration_number, message_type, content, - token_count, model_used, temperature, response_time_ms, metadata - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) - RETURNING id - """, - message.task_id, - message.iteration_number, - message.message_type.value, - message.content, - message.token_count, - message.model_used, - message.temperature, - message.response_time_ms, - json.dumps(message.metadata) if message.metadata else {}, - ) - - return str(message_id) +""" +Conversation Manager for FuzeAgent Claude Code Integration + +Manages and stores complete conversations between agents and Claude Code, +providing comprehensive audit trails, debugging capabilities, and learning data. +""" + +import asyncio +import json +import logging +import time +import uuid +from dataclasses import dataclass +from datetime import datetime +from enum import Enum +from typing import Any, Dict, List, Optional + +# Importable both as `services.orchestrator.conversation_manager` (relative) and +# flat with services/orchestrator on sys.path (as the tests import it). +try: + from .database import get_db_connection +except ImportError: # pragma: no cover - flat import (no parent package) + from database import get_db_connection + +logger = logging.getLogger(__name__) + + +class MessageType(str, Enum): + USER_PROMPT = "user_prompt" + CLAUDE_RESPONSE = "claude_response" + SYSTEM_MESSAGE = "system_message" + ERROR_MESSAGE = "error_message" + CODE_EXECUTION = "code_execution" + TEST_RESULT = "test_result" + + +class InteractionType(str, Enum): + QUESTION = "question" + CLARIFICATION = "clarification" + APPROVAL_REQUEST = "approval_request" + ERROR_REPORT = "error_report" + PROGRESS_UPDATE = "progress_update" + + +@dataclass +class ConversationMessage: + """Represents a single message in a Claude Code conversation""" + + task_id: str + iteration_number: int + message_type: MessageType + content: str + token_count: Optional[int] = None + model_used: Optional[str] = None + temperature: Optional[float] = None + response_time_ms: Optional[int] = None + metadata: Optional[Dict[str, Any]] = None + + +@dataclass +class ConversationSession: + """Represents a complete conversation session""" + + agent_id: str + task_id: str + sandbox_id: str + session_started_at: datetime + session_ended_at: Optional[datetime] = None + total_messages: int = 0 + total_tokens: int = 0 + status: str = "active" + metadata: Optional[Dict[str, Any]] = None + + +class ConversationManager: + """ + Manages Claude Code conversations and provides comprehensive tracking. + + Features: + - Full conversation storage and retrieval + - Token usage tracking and cost analysis + - Human interaction management + - Code generation tracking + - Performance metrics collection + """ + + def __init__(self): + self.active_sessions: Dict[str, ConversationSession] = {} + + async def start_conversation_session( + self, + agent_id: str, + task_id: str, + sandbox_id: str, + metadata: Optional[Dict[str, Any]] = None, + ) -> str: + """Start a new conversation session for an agent""" + + session = ConversationSession( + agent_id=agent_id, + task_id=task_id, + sandbox_id=sandbox_id, + session_started_at=datetime.now(), + metadata=metadata or {}, + ) + + # Store session in database + session_id = await self._store_conversation_session(session) + session.metadata = session.metadata or {} + session.metadata["session_id"] = session_id + + # Track active session + self.active_sessions[session_id] = session + + logger.info( + f"Started conversation session {session_id} for agent {agent_id}, task {task_id}" + ) + return session_id + + async def end_conversation_session(self, session_id: str) -> bool: + """End a conversation session""" + + session = self.active_sessions.get(session_id) + if not session: + logger.warning(f"Session {session_id} not found in active sessions") + return False + + session.session_ended_at = datetime.now() + session.status = "completed" + + # Update database + await self._update_conversation_session(session_id, session) + + # Remove from active sessions + del self.active_sessions[session_id] + + logger.info(f"Ended conversation session {session_id}") + return True + + async def store_message( + self, + session_id: str, + message: ConversationMessage, + start_time: Optional[float] = None, + ) -> str: + """Store a conversation message""" + + # Calculate response time if start_time provided + if start_time and message.message_type == MessageType.CLAUDE_RESPONSE: + message.response_time_ms = int((time.time() - start_time) * 1000) + + # Store message in database + message_id = await self._store_claude_conversation(message) + + # Update session statistics + session = self.active_sessions.get(session_id) + if session: + session.total_messages += 1 + if message.token_count: + session.total_tokens += message.token_count + + logger.debug(f"Stored message {message_id} for session {session_id}") + return message_id + + async def store_user_prompt( + self, + session_id: str, + task_id: str, + iteration_number: int, + prompt: str, + model: str = "claude-3-5-sonnet-20241022", + temperature: float = 0.3, + metadata: Optional[Dict[str, Any]] = None, + ) -> str: + """Store a user prompt to Claude""" + + message = ConversationMessage( + task_id=task_id, + iteration_number=iteration_number, + message_type=MessageType.USER_PROMPT, + content=prompt, + model_used=model, + temperature=temperature, + metadata=metadata or {}, + ) + + return await self.store_message(session_id, message) + + async def store_claude_response( + self, + session_id: str, + task_id: str, + iteration_number: int, + response: str, + token_count: Optional[int] = None, + model: str = "claude-3-5-sonnet-20241022", + start_time: Optional[float] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> str: + """Store Claude's response""" + + message = ConversationMessage( + task_id=task_id, + iteration_number=iteration_number, + message_type=MessageType.CLAUDE_RESPONSE, + content=response, + token_count=token_count, + model_used=model, + metadata=metadata or {}, + ) + + return await self.store_message(session_id, message, start_time) + + async def store_code_generation( + self, + task_id: str, + iteration_number: int, + file_path: str, + file_type: str, + language: str, + content: str, + commit_hash: Optional[str] = None, + test_results: Optional[Dict[str, Any]] = None, + quality_metrics: Optional[Dict[str, Any]] = None, + ) -> str: + """Store generated code with metadata""" + + async with get_db_connection() as conn: + code_id = await conn.fetchval( + """ + INSERT INTO code_generations ( + task_id, iteration_number, file_path, file_type, language, + content, commit_hash, test_results, quality_metrics + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) + RETURNING id + """, + task_id, + iteration_number, + file_path, + file_type, + language, + content, + commit_hash, + json.dumps(test_results) if test_results else None, + json.dumps(quality_metrics) if quality_metrics else None, + ) + + logger.info(f"Stored code generation {code_id} for task {task_id}") + return str(code_id) + + async def store_human_interaction( + self, + task_id: str, + iteration_number: int, + interaction_type: InteractionType, + agent_message: str, + human_response: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> str: + """Store human-in-the-loop interaction""" + + async with get_db_connection() as conn: + interaction_id = await conn.fetchval( + """ + INSERT INTO human_interactions ( + task_id, iteration_number, interaction_type, + agent_message, human_response, metadata + ) VALUES ($1, $2, $3, $4, $5, $6) + RETURNING id + """, + task_id, + iteration_number, + interaction_type.value, + agent_message, + human_response, + json.dumps(metadata) if metadata else {}, + ) + + logger.info(f"Stored human interaction {interaction_id} for task {task_id}") + return str(interaction_id) + + async def update_human_response( + self, interaction_id: str, human_response: str + ) -> bool: + """Update human response to an interaction""" + + async with get_db_connection() as conn: + result = await conn.execute( + """ + UPDATE human_interactions + SET human_response = $1, + responded_at = CURRENT_TIMESTAMP, + response_time_seconds = EXTRACT(EPOCH FROM (CURRENT_TIMESTAMP - asked_at)) + WHERE id = $2 + """, + human_response, + interaction_id, + ) + + success = result != "UPDATE 0" + if success: + logger.info(f"Updated human response for interaction {interaction_id}") + return success + + async def store_performance_metric( + self, + agent_id: str, + task_id: str, + metric_type: str, + metric_value: float, + metric_unit: Optional[str] = None, + context: Optional[Dict[str, Any]] = None, + ) -> str: + """Store agent performance metric""" + + async with get_db_connection() as conn: + metric_id = await conn.fetchval( + """ + INSERT INTO agent_performance_metrics ( + agent_id, task_id, metric_type, metric_value, + metric_unit, context + ) VALUES ($1, $2, $3, $4, $5, $6) + RETURNING id + """, + agent_id, + task_id, + metric_type, + metric_value, + metric_unit, + json.dumps(context) if context else {}, + ) + + logger.debug( + f"Stored performance metric {metric_id}: {metric_type}={metric_value}" + ) + return str(metric_id) + + async def get_conversation_history( + self, + task_id: str, + iteration_number: Optional[int] = None, + message_types: Optional[List[MessageType]] = None, + limit: Optional[int] = None, + ) -> List[Dict[str, Any]]: + """Get conversation history for a task""" + + conditions = ["task_id = $1"] + params = [task_id] + param_count = 1 + + if iteration_number is not None: + param_count += 1 + conditions.append(f"iteration_number = ${param_count}") + params.append(iteration_number) + + if message_types: + param_count += 1 + conditions.append(f"message_type = ANY(${param_count})") + params.append([mt.value for mt in message_types]) + + where_clause = " AND ".join(conditions) + limit_clause = "" + if limit: + param_count += 1 + limit_clause = f"LIMIT ${param_count}" + params.append(limit) + + async with get_db_connection() as conn: + rows = await conn.fetch( + f""" + SELECT * FROM claude_conversations + WHERE {where_clause} + ORDER BY created_at ASC + {limit_clause} + """, # nosec B608 -- where/limit clauses are fixed fragments with $N placeholders; all values bound as query params + *params, + ) + + return [dict(row) for row in rows] + + async def get_conversation_summary(self, task_id: str) -> Dict[str, Any]: + """Get conversation summary with statistics""" + + async with get_db_connection() as conn: + # Get message statistics + stats = await conn.fetchrow( + """ + SELECT + COUNT(*) as total_messages, + SUM(token_count) as total_tokens, + AVG(response_time_ms) as avg_response_time, + MAX(iteration_number) as max_iteration + FROM claude_conversations + WHERE task_id = $1 + """, + task_id, + ) + + # Get message type breakdown + type_breakdown = await conn.fetch( + """ + SELECT message_type, COUNT(*) as count + FROM claude_conversations + WHERE task_id = $1 + GROUP BY message_type + """, + task_id, + ) + + # Get human interactions + human_interactions = await conn.fetchrow( + """ + SELECT + COUNT(*) as total_interactions, + COUNT(human_response) as responded_interactions, + AVG(response_time_seconds) as avg_response_time + FROM human_interactions + WHERE task_id = $1 + """, + task_id, + ) + + return { + "task_id": task_id, + "total_messages": stats["total_messages"] or 0, + "total_tokens": stats["total_tokens"] or 0, + "avg_response_time_ms": float(stats["avg_response_time"] or 0), + "max_iteration": stats["max_iteration"] or 0, + "message_types": { + row["message_type"]: row["count"] for row in type_breakdown + }, + "human_interactions": { + "total": human_interactions["total_interactions"] or 0, + "responded": human_interactions["responded_interactions"] or 0, + "avg_response_time_seconds": float( + human_interactions["avg_response_time"] or 0 + ), + }, + } + + async def get_code_generations( + self, + task_id: str, + iteration_number: Optional[int] = None, + file_type: Optional[str] = None, + language: Optional[str] = None, + ) -> List[Dict[str, Any]]: + """Get code generations for a task""" + + conditions = ["task_id = $1"] + params = [task_id] + param_count = 1 + + if iteration_number is not None: + param_count += 1 + conditions.append(f"iteration_number = ${param_count}") + params.append(iteration_number) + + if file_type: + param_count += 1 + conditions.append(f"file_type = ${param_count}") + params.append(file_type) + + if language: + param_count += 1 + conditions.append(f"language = ${param_count}") + params.append(language) + + where_clause = " AND ".join(conditions) + + async with get_db_connection() as conn: + rows = await conn.fetch( + f""" + SELECT * FROM code_generations + WHERE {where_clause} + ORDER BY generated_at ASC + """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params + *params, + ) + + return [dict(row) for row in rows] + + async def get_agent_performance_metrics( + self, + agent_id: Optional[str] = None, + task_id: Optional[str] = None, + metric_types: Optional[List[str]] = None, + time_range_hours: Optional[int] = None, + ) -> List[Dict[str, Any]]: + """Get agent performance metrics""" + + conditions = [] + params = [] + param_count = 0 + + if agent_id: + param_count += 1 + conditions.append(f"agent_id = ${param_count}") + params.append(agent_id) + + if task_id: + param_count += 1 + conditions.append(f"task_id = ${param_count}") + params.append(task_id) + + if metric_types: + param_count += 1 + conditions.append(f"metric_type = ANY(${param_count})") + params.append(metric_types) + + if time_range_hours: + param_count += 1 + conditions.append( + f"measured_at >= NOW() - (INTERVAL '1 hour' * ${param_count})" + ) + params.append(time_range_hours) + + where_clause = "WHERE " + " AND ".join(conditions) if conditions else "" + + async with get_db_connection() as conn: + rows = await conn.fetch( + f""" + SELECT * FROM agent_performance_metrics + {where_clause} + ORDER BY measured_at DESC + """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values (incl. interval multiplier) bound as query params + *params, + ) + + return [dict(row) for row in rows] + + # Private methods + + async def _store_conversation_session(self, session: ConversationSession) -> str: + """Store conversation session in database""" + + async with get_db_connection() as conn: + session_id = await conn.fetchval( + """ + INSERT INTO agent_conversation_sessions ( + agent_id, task_id, sandbox_id, session_started_at, + total_messages, total_tokens, status, metadata + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + RETURNING id + """, + session.agent_id, + session.task_id, + session.sandbox_id, + session.session_started_at, + session.total_messages, + session.total_tokens, + session.status, + json.dumps(session.metadata) if session.metadata else {}, + ) + + return str(session_id) + + async def _update_conversation_session( + self, session_id: str, session: ConversationSession + ): + """Update conversation session in database""" + + async with get_db_connection() as conn: + await conn.execute( + """ + UPDATE agent_conversation_sessions + SET session_ended_at = $1, total_messages = $2, + total_tokens = $3, status = $4, metadata = $5 + WHERE id = $6 + """, + session.session_ended_at, + session.total_messages, + session.total_tokens, + session.status, + json.dumps(session.metadata) if session.metadata else {}, + session_id, + ) + + async def _store_claude_conversation(self, message: ConversationMessage) -> str: + """Store Claude conversation message in database""" + + async with get_db_connection() as conn: + message_id = await conn.fetchval( + """ + INSERT INTO claude_conversations ( + task_id, iteration_number, message_type, content, + token_count, model_used, temperature, response_time_ms, metadata + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) + RETURNING id + """, + message.task_id, + message.iteration_number, + message.message_type.value, + message.content, + message.token_count, + message.model_used, + message.temperature, + message.response_time_ms, + json.dumps(message.metadata) if message.metadata else {}, + ) + + return str(message_id) diff --git a/services/orchestrator/coordination_endpoints.py b/services/orchestrator/coordination_endpoints.py index ee08b34..cdcd244 100644 --- a/services/orchestrator/coordination_endpoints.py +++ b/services/orchestrator/coordination_endpoints.py @@ -1,625 +1,625 @@ -""" -Cross-Product Coordination API Endpoints - -This module provides REST API endpoints for managing cross-product coordination -within the WCG ecosystem. Enables centralized coordination between FuzeAgent, -FuzeFront, HubHit, DeployAI, and other WCG products. -""" - -import logging -from datetime import date, datetime -from typing import Any, Dict, List, Optional - -import asyncpg -from fastapi import APIRouter, Depends, HTTPException, Path, Query -from pydantic import BaseModel, Field - -from .database import get_db_connection - -logger = logging.getLogger(__name__) -router = APIRouter(prefix="/coordination", tags=["Cross-Product Coordination"]) - - -# Pydantic Models -class ProductRegistration(BaseModel): - id: str = Field(..., description="Unique product identifier") - name: str = Field(..., description="Product display name") - version: str = Field(..., description="Current product version") - endpoints: List[str] = Field(default=[], description="API endpoints exposed") - dependencies: List[str] = Field(default=[], description="Product dependencies") - resource_requirements: Dict[str, Any] = Field( - default={}, description="Resource needs" - ) - team_contacts: List[str] = Field(default=[], description="Team contact information") - priority_level: int = Field(default=5, ge=1, le=10, description="Business priority") - metadata: Dict[str, Any] = Field(default={}, description="Additional metadata") - - -class CoordinationRequestCreate(BaseModel): - requesting_product: str = Field(..., description="Product making the request") - target_products: List[str] = Field( - ..., description="Target products for coordination" - ) - coordination_type: str = Field(..., description="Type of coordination needed") - scope: str = Field(default="product_group", description="Coordination scope") - priority: str = Field(default="medium", description="Request priority") - title: str = Field(..., description="Brief title for the request") - description: str = Field(..., description="Detailed description") - resource_requirements: Dict[str, Any] = Field( - default={}, description="Required resources" - ) - proposed_timeline: Dict[str, Any] = Field( - default={}, description="Proposed timeline" - ) - - -class CoordinationResolution(BaseModel): - resolution_plan: Dict[str, Any] = Field(..., description="Detailed resolution plan") - resolver_id: str = Field(..., description="ID of the resolver") - notes: Optional[str] = Field(None, description="Additional resolution notes") - - -class ResourceAllocationCreate(BaseModel): - product_id: str = Field(..., description="Product requesting allocation") - resource_type: str = Field(..., description="Type of resource") - resource_name: str = Field(..., description="Specific resource name") - allocation_details: Dict[str, Any] = Field( - default={}, description="Allocation specifics" - ) - valid_until: Optional[datetime] = Field(None, description="Allocation expiry") - - -# Product Registration Endpoints -@router.post( - "/products/register", summary="Register a new product in coordination system" -) -async def register_product(product: ProductRegistration): - """Register a new product in the cross-product coordination system""" - try: - async with get_db_connection() as conn: - await conn.execute( - """ - INSERT INTO product_registry ( - id, name, version, endpoints, dependencies, - resource_requirements, team_contacts, priority_level, - metadata, registered_at, updated_at - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) - ON CONFLICT (id) DO UPDATE SET - name = $2, version = $3, endpoints = $4, dependencies = $5, - resource_requirements = $6, team_contacts = $7, - priority_level = $8, metadata = $9, updated_at = $11 - """, - product.id, - product.name, - product.version, - product.endpoints, - product.dependencies, - product.resource_requirements, - product.team_contacts, - product.priority_level, - product.metadata, - datetime.utcnow(), - datetime.utcnow(), - ) - - return { - "status": "success", - "product_id": product.id, - "message": "Product registered successfully", - } - - except Exception as e: - logger.error(f"Error registering product {product.id}: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to register product: {str(e)}" - ) - - -@router.get("/products", summary="List all registered products") -async def list_products( - priority_min: int = Query(1, ge=1, le=10, description="Minimum priority level"), - active_only: bool = Query(True, description="Show only active products"), -): - """List all products registered in the coordination system""" - try: - async with get_db_connection() as conn: - products = await conn.fetch( - """ - SELECT id, name, version, priority_level, metadata, - endpoints, dependencies, registered_at, updated_at - FROM product_registry - WHERE priority_level >= $1 - ORDER BY priority_level DESC, name ASC - """, - priority_min, - ) - - return { - "products": [dict(product) for product in products], - "total_count": len(products), - } - - except Exception as e: - logger.error(f"Error listing products: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to list products: {str(e)}" - ) - - -@router.get("/products/{product_id}", summary="Get specific product details") -async def get_product(product_id: str = Path(..., description="Product ID")): - """Get detailed information about a specific product""" - try: - async with get_db_connection() as conn: - product = await conn.fetchrow( - """ - SELECT * FROM product_registry WHERE id = $1 - """, - product_id, - ) - - if not product: - raise HTTPException( - status_code=404, detail=f"Product {product_id} not found" - ) - - # Get active coordination requests involving this product - async with get_db_connection() as conn: - coordination_requests = await conn.fetch( - """ - SELECT id, title, coordination_type, priority, status, created_at - FROM coordination_requests - WHERE requesting_product = $1 - OR $1 = ANY(string_to_array(replace(replace(target_products::text, '[', ''), ']', ''), ',')) - ORDER BY created_at DESC LIMIT 10 - """, - product_id, - ) - - return { - "product": dict(product), - "active_coordination_requests": [ - dict(req) for req in coordination_requests - ], - } - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting product {product_id}: {e}") - raise HTTPException(status_code=500, detail=f"Failed to get product: {str(e)}") - - -# Coordination Request Endpoints -@router.post("/requests", summary="Create a new coordination request") -async def create_coordination_request(request: CoordinationRequestCreate): - """Create a new cross-product coordination request""" - try: - # Validate requesting product exists - async with get_db_connection() as conn: - requesting_product = await conn.fetchrow( - """ - SELECT id FROM product_registry WHERE id = $1 - """, - request.requesting_product, - ) - - if not requesting_product: - raise HTTPException( - status_code=400, - detail=f"Requesting product {request.requesting_product} not found", - ) - - # Create coordination request - async with get_db_connection() as conn: - request_id = await conn.fetchval( - """ - INSERT INTO coordination_requests ( - requesting_product, target_products, coordination_type, - scope, priority, title, description, resource_requirements, - proposed_timeline, stakeholders, created_at, updated_at - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) - RETURNING id - """, - request.requesting_product, - request.target_products, - request.coordination_type, - request.scope, - request.priority, - request.title, - request.description, - request.resource_requirements, - request.proposed_timeline, - [], # stakeholders - could be auto-populated - datetime.utcnow(), - datetime.utcnow(), - ) - - # Log coordination history - async with get_db_connection() as conn: - await conn.execute( - """ - INSERT INTO coordination_history ( - coordination_request_id, action, actor_type, details - ) VALUES ($1, $2, $3, $4) - """, - request_id, - "created", - "system", - {"created_via": "api"}, - ) - - return { - "status": "success", - "request_id": str(request_id), - "message": "Coordination request created successfully", - } - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error creating coordination request: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to create coordination request: {str(e)}" - ) - - -@router.get("/requests", summary="List coordination requests") -async def list_coordination_requests( - status: Optional[str] = Query(None, description="Filter by status"), - priority: Optional[str] = Query(None, description="Filter by priority"), - product_id: Optional[str] = Query(None, description="Filter by product"), - limit: int = Query(50, ge=1, le=200, description="Max number of results"), -): - """List coordination requests with optional filters""" - try: - where_conditions = [] - params = [] - param_count = 0 - - if status: - param_count += 1 - where_conditions.append(f"status = ${param_count}") - params.append(status) - - if priority: - param_count += 1 - where_conditions.append(f"priority = ${param_count}") - params.append(priority) - - if product_id: - param_count += 1 - where_conditions.append( - f"(requesting_product = ${param_count} OR ${param_count} = ANY(string_to_array(replace(replace(target_products::text, '[', ''), ']', ''), ',')))" - ) - params.append(product_id) - - where_clause = ( - " WHERE " + " AND ".join(where_conditions) if where_conditions else "" - ) - param_count += 1 - params.append(limit) - - query = f""" - SELECT id, requesting_product, target_products, coordination_type, - scope, priority, status, title, description, created_at, updated_at - FROM coordination_requests - {where_clause} - ORDER BY - CASE priority - WHEN 'critical' THEN 1 - WHEN 'high' THEN 2 - WHEN 'medium' THEN 3 - WHEN 'low' THEN 4 - END, - created_at DESC - LIMIT ${param_count} - """ # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params - - async with get_db_connection() as conn: - requests = await conn.fetch(query, *params) - - return { - "coordination_requests": [dict(req) for req in requests], - "total_count": len(requests), - } - - except Exception as e: - logger.error(f"Error listing coordination requests: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to list coordination requests: {str(e)}" - ) - - -@router.get("/requests/{request_id}", summary="Get coordination request details") -async def get_coordination_request( - request_id: str = Path(..., description="Coordination request ID") -): - """Get detailed information about a specific coordination request""" - try: - async with get_db_connection() as conn: - request = await conn.fetchrow( - """ - SELECT * FROM coordination_requests WHERE id = $1 - """, - request_id, - ) - - if not request: - raise HTTPException( - status_code=404, detail=f"Coordination request {request_id} not found" - ) - - # Get coordination history - async with get_db_connection() as conn: - history = await conn.fetch( - """ - SELECT action, actor_id, actor_type, details, timestamp - FROM coordination_history - WHERE coordination_request_id = $1 - ORDER BY timestamp ASC - """, - request_id, - ) - - return { - "coordination_request": dict(request), - "history": [dict(h) for h in history], - } - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting coordination request {request_id}: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to get coordination request: {str(e)}" - ) - - -@router.put("/requests/{request_id}/resolve", summary="Resolve coordination request") -async def resolve_coordination_request( - request_id: str = Path(..., description="Coordination request ID"), - resolution: CoordinationResolution = ..., -): - """Resolve a coordination request with a specific plan""" - try: - async with get_db_connection() as conn: - # Check if request exists and is pending - existing_request = await conn.fetchrow( - """ - SELECT id, status FROM coordination_requests WHERE id = $1 - """, - request_id, - ) - - if not existing_request: - raise HTTPException( - status_code=404, detail=f"Coordination request {request_id} not found" - ) - - if existing_request["status"] not in ["pending", "in_progress"]: - raise HTTPException( - status_code=400, - detail=f"Request is already {existing_request['status']}", - ) - - # Update request status and resolution - async with get_db_connection() as conn: - await conn.execute( - """ - UPDATE coordination_requests SET - status = 'resolved', - resolution_plan = $2, - resolved_at = $3, - updated_at = $4 - WHERE id = $1 - """, - request_id, - resolution.resolution_plan, - datetime.utcnow(), - datetime.utcnow(), - ) - - # Log resolution in history - async with get_db_connection() as conn: - await conn.execute( - """ - INSERT INTO coordination_history ( - coordination_request_id, action, actor_id, actor_type, details - ) VALUES ($1, $2, $3, $4, $5) - """, - request_id, - "resolved", - resolution.resolver_id, - "agent", - { - "resolution_plan": resolution.resolution_plan, - "notes": resolution.notes, - }, - ) - - return { - "status": "success", - "request_id": request_id, - "message": "Coordination request resolved successfully", - } - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error resolving coordination request {request_id}: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to resolve coordination request: {str(e)}" - ) - - -# Resource Management Endpoints -@router.post("/resources/allocate", summary="Allocate resources to a product") -async def allocate_resource(allocation: ResourceAllocationCreate): - """Allocate a resource to a specific product""" - try: - async with get_db_connection() as conn: - allocation_id = await conn.fetchval( - """ - INSERT INTO resource_allocations ( - product_id, resource_type, resource_name, allocation_details, - valid_until, status, created_at, updated_at - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) - RETURNING id - """, - allocation.product_id, - allocation.resource_type, - allocation.resource_name, - allocation.allocation_details, - allocation.valid_until, - "active", - datetime.utcnow(), - datetime.utcnow(), - ) - - return { - "status": "success", - "allocation_id": str(allocation_id), - "message": "Resource allocated successfully", - } - - except asyncpg.UniqueViolationError: - raise HTTPException( - status_code=409, - detail=f"Resource {allocation.resource_name} of type {allocation.resource_type} is already allocated", - ) - except Exception as e: - logger.error(f"Error allocating resource: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to allocate resource: {str(e)}" - ) - - -@router.get("/resources", summary="List resource allocations") -async def list_resource_allocations( - product_id: Optional[str] = Query(None, description="Filter by product"), - resource_type: Optional[str] = Query(None, description="Filter by resource type"), - status: Optional[str] = Query(None, description="Filter by status"), -): - """List current resource allocations""" - try: - where_conditions = [] - params = [] - param_count = 0 - - if product_id: - param_count += 1 - where_conditions.append(f"product_id = ${param_count}") - params.append(product_id) - - if resource_type: - param_count += 1 - where_conditions.append(f"resource_type = ${param_count}") - params.append(resource_type) - - if status: - param_count += 1 - where_conditions.append(f"status = ${param_count}") - params.append(status) - - where_clause = ( - " WHERE " + " AND ".join(where_conditions) if where_conditions else "" - ) - - query = f""" - SELECT ra.*, pr.name as product_name - FROM resource_allocations ra - LEFT JOIN product_registry pr ON ra.product_id = pr.id - {where_clause} - ORDER BY ra.created_at DESC - """ # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params - - async with get_db_connection() as conn: - allocations = await conn.fetch(query, *params) - - return { - "resource_allocations": [dict(alloc) for alloc in allocations], - "total_count": len(allocations), - } - - except Exception as e: - logger.error(f"Error listing resource allocations: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to list resource allocations: {str(e)}" - ) - - -# Dashboard and Status Endpoints -@router.get("/status", summary="Get overall coordination system status") -async def get_coordination_status(): - """Get overall status of the cross-product coordination system""" - try: - async with get_db_connection() as conn: - # Get request counts by status - request_stats = await conn.fetch(""" - SELECT status, priority, COUNT(*) as count - FROM coordination_requests - GROUP BY status, priority - ORDER BY status, priority - """) - - # Get product count - product_count = await conn.fetchval(""" - SELECT COUNT(*) FROM product_registry - """) - - # Get resource allocation stats - resource_stats = await conn.fetch(""" - SELECT resource_type, status, COUNT(*) as count - FROM resource_allocations - GROUP BY resource_type, status - """) - - # Get recent activity - recent_activity = await conn.fetch(""" - SELECT ch.action, ch.timestamp, cr.title, cr.requesting_product - FROM coordination_history ch - JOIN coordination_requests cr ON ch.coordination_request_id = cr.id - ORDER BY ch.timestamp DESC - LIMIT 10 - """) - - return { - "system_status": "operational", - "registered_products": product_count, - "coordination_request_stats": [dict(stat) for stat in request_stats], - "resource_allocation_stats": [dict(stat) for stat in resource_stats], - "recent_activity": [dict(activity) for activity in recent_activity], - "last_updated": datetime.utcnow().isoformat(), - } - - except Exception as e: - logger.error(f"Error getting coordination status: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to get coordination status: {str(e)}" - ) - - -@router.get("/protocols", summary="List coordination protocols") -async def list_coordination_protocols(): - """List available coordination protocols and procedures""" - try: - async with get_db_connection() as conn: - protocols = await conn.fetch(""" - SELECT protocol_name, coordination_type, scope, procedure_steps, - required_approvals, sla_requirements, is_active, version - FROM coordination_protocols - WHERE is_active = true - ORDER BY protocol_name ASC - """) - - return { - "coordination_protocols": [dict(protocol) for protocol in protocols], - "total_count": len(protocols), - } - - except Exception as e: - logger.error(f"Error listing coordination protocols: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to list coordination protocols: {str(e)}" - ) +""" +Cross-Product Coordination API Endpoints + +This module provides REST API endpoints for managing cross-product coordination +within the WCG ecosystem. Enables centralized coordination between FuzeAgent, +FuzeFront, HubHit, DeployAI, and other WCG products. +""" + +import logging +from datetime import date, datetime +from typing import Any, Dict, List, Optional + +import asyncpg +from fastapi import APIRouter, Depends, HTTPException, Path, Query +from pydantic import BaseModel, Field + +from .database import get_db_connection + +logger = logging.getLogger(__name__) +router = APIRouter(prefix="/coordination", tags=["Cross-Product Coordination"]) + + +# Pydantic Models +class ProductRegistration(BaseModel): + id: str = Field(..., description="Unique product identifier") + name: str = Field(..., description="Product display name") + version: str = Field(..., description="Current product version") + endpoints: List[str] = Field(default=[], description="API endpoints exposed") + dependencies: List[str] = Field(default=[], description="Product dependencies") + resource_requirements: Dict[str, Any] = Field( + default={}, description="Resource needs" + ) + team_contacts: List[str] = Field(default=[], description="Team contact information") + priority_level: int = Field(default=5, ge=1, le=10, description="Business priority") + metadata: Dict[str, Any] = Field(default={}, description="Additional metadata") + + +class CoordinationRequestCreate(BaseModel): + requesting_product: str = Field(..., description="Product making the request") + target_products: List[str] = Field( + ..., description="Target products for coordination" + ) + coordination_type: str = Field(..., description="Type of coordination needed") + scope: str = Field(default="product_group", description="Coordination scope") + priority: str = Field(default="medium", description="Request priority") + title: str = Field(..., description="Brief title for the request") + description: str = Field(..., description="Detailed description") + resource_requirements: Dict[str, Any] = Field( + default={}, description="Required resources" + ) + proposed_timeline: Dict[str, Any] = Field( + default={}, description="Proposed timeline" + ) + + +class CoordinationResolution(BaseModel): + resolution_plan: Dict[str, Any] = Field(..., description="Detailed resolution plan") + resolver_id: str = Field(..., description="ID of the resolver") + notes: Optional[str] = Field(None, description="Additional resolution notes") + + +class ResourceAllocationCreate(BaseModel): + product_id: str = Field(..., description="Product requesting allocation") + resource_type: str = Field(..., description="Type of resource") + resource_name: str = Field(..., description="Specific resource name") + allocation_details: Dict[str, Any] = Field( + default={}, description="Allocation specifics" + ) + valid_until: Optional[datetime] = Field(None, description="Allocation expiry") + + +# Product Registration Endpoints +@router.post( + "/products/register", summary="Register a new product in coordination system" +) +async def register_product(product: ProductRegistration): + """Register a new product in the cross-product coordination system""" + try: + async with get_db_connection() as conn: + await conn.execute( + """ + INSERT INTO product_registry ( + id, name, version, endpoints, dependencies, + resource_requirements, team_contacts, priority_level, + metadata, registered_at, updated_at + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) + ON CONFLICT (id) DO UPDATE SET + name = $2, version = $3, endpoints = $4, dependencies = $5, + resource_requirements = $6, team_contacts = $7, + priority_level = $8, metadata = $9, updated_at = $11 + """, + product.id, + product.name, + product.version, + product.endpoints, + product.dependencies, + product.resource_requirements, + product.team_contacts, + product.priority_level, + product.metadata, + datetime.utcnow(), + datetime.utcnow(), + ) + + return { + "status": "success", + "product_id": product.id, + "message": "Product registered successfully", + } + + except Exception as e: + logger.error(f"Error registering product {product.id}: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to register product: {str(e)}" + ) + + +@router.get("/products", summary="List all registered products") +async def list_products( + priority_min: int = Query(1, ge=1, le=10, description="Minimum priority level"), + active_only: bool = Query(True, description="Show only active products"), +): + """List all products registered in the coordination system""" + try: + async with get_db_connection() as conn: + products = await conn.fetch( + """ + SELECT id, name, version, priority_level, metadata, + endpoints, dependencies, registered_at, updated_at + FROM product_registry + WHERE priority_level >= $1 + ORDER BY priority_level DESC, name ASC + """, + priority_min, + ) + + return { + "products": [dict(product) for product in products], + "total_count": len(products), + } + + except Exception as e: + logger.error(f"Error listing products: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to list products: {str(e)}" + ) + + +@router.get("/products/{product_id}", summary="Get specific product details") +async def get_product(product_id: str = Path(..., description="Product ID")): + """Get detailed information about a specific product""" + try: + async with get_db_connection() as conn: + product = await conn.fetchrow( + """ + SELECT * FROM product_registry WHERE id = $1 + """, + product_id, + ) + + if not product: + raise HTTPException( + status_code=404, detail=f"Product {product_id} not found" + ) + + # Get active coordination requests involving this product + async with get_db_connection() as conn: + coordination_requests = await conn.fetch( + """ + SELECT id, title, coordination_type, priority, status, created_at + FROM coordination_requests + WHERE requesting_product = $1 + OR $1 = ANY(string_to_array(replace(replace(target_products::text, '[', ''), ']', ''), ',')) + ORDER BY created_at DESC LIMIT 10 + """, + product_id, + ) + + return { + "product": dict(product), + "active_coordination_requests": [ + dict(req) for req in coordination_requests + ], + } + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting product {product_id}: {e}") + raise HTTPException(status_code=500, detail=f"Failed to get product: {str(e)}") + + +# Coordination Request Endpoints +@router.post("/requests", summary="Create a new coordination request") +async def create_coordination_request(request: CoordinationRequestCreate): + """Create a new cross-product coordination request""" + try: + # Validate requesting product exists + async with get_db_connection() as conn: + requesting_product = await conn.fetchrow( + """ + SELECT id FROM product_registry WHERE id = $1 + """, + request.requesting_product, + ) + + if not requesting_product: + raise HTTPException( + status_code=400, + detail=f"Requesting product {request.requesting_product} not found", + ) + + # Create coordination request + async with get_db_connection() as conn: + request_id = await conn.fetchval( + """ + INSERT INTO coordination_requests ( + requesting_product, target_products, coordination_type, + scope, priority, title, description, resource_requirements, + proposed_timeline, stakeholders, created_at, updated_at + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) + RETURNING id + """, + request.requesting_product, + request.target_products, + request.coordination_type, + request.scope, + request.priority, + request.title, + request.description, + request.resource_requirements, + request.proposed_timeline, + [], # stakeholders - could be auto-populated + datetime.utcnow(), + datetime.utcnow(), + ) + + # Log coordination history + async with get_db_connection() as conn: + await conn.execute( + """ + INSERT INTO coordination_history ( + coordination_request_id, action, actor_type, details + ) VALUES ($1, $2, $3, $4) + """, + request_id, + "created", + "system", + {"created_via": "api"}, + ) + + return { + "status": "success", + "request_id": str(request_id), + "message": "Coordination request created successfully", + } + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error creating coordination request: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to create coordination request: {str(e)}" + ) + + +@router.get("/requests", summary="List coordination requests") +async def list_coordination_requests( + status: Optional[str] = Query(None, description="Filter by status"), + priority: Optional[str] = Query(None, description="Filter by priority"), + product_id: Optional[str] = Query(None, description="Filter by product"), + limit: int = Query(50, ge=1, le=200, description="Max number of results"), +): + """List coordination requests with optional filters""" + try: + where_conditions = [] + params = [] + param_count = 0 + + if status: + param_count += 1 + where_conditions.append(f"status = ${param_count}") + params.append(status) + + if priority: + param_count += 1 + where_conditions.append(f"priority = ${param_count}") + params.append(priority) + + if product_id: + param_count += 1 + where_conditions.append( + f"(requesting_product = ${param_count} OR ${param_count} = ANY(string_to_array(replace(replace(target_products::text, '[', ''), ']', ''), ',')))" + ) + params.append(product_id) + + where_clause = ( + " WHERE " + " AND ".join(where_conditions) if where_conditions else "" + ) + param_count += 1 + params.append(limit) + + query = f""" + SELECT id, requesting_product, target_products, coordination_type, + scope, priority, status, title, description, created_at, updated_at + FROM coordination_requests + {where_clause} + ORDER BY + CASE priority + WHEN 'critical' THEN 1 + WHEN 'high' THEN 2 + WHEN 'medium' THEN 3 + WHEN 'low' THEN 4 + END, + created_at DESC + LIMIT ${param_count} + """ # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params + + async with get_db_connection() as conn: + requests = await conn.fetch(query, *params) + + return { + "coordination_requests": [dict(req) for req in requests], + "total_count": len(requests), + } + + except Exception as e: + logger.error(f"Error listing coordination requests: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to list coordination requests: {str(e)}" + ) + + +@router.get("/requests/{request_id}", summary="Get coordination request details") +async def get_coordination_request( + request_id: str = Path(..., description="Coordination request ID") +): + """Get detailed information about a specific coordination request""" + try: + async with get_db_connection() as conn: + request = await conn.fetchrow( + """ + SELECT * FROM coordination_requests WHERE id = $1 + """, + request_id, + ) + + if not request: + raise HTTPException( + status_code=404, detail=f"Coordination request {request_id} not found" + ) + + # Get coordination history + async with get_db_connection() as conn: + history = await conn.fetch( + """ + SELECT action, actor_id, actor_type, details, timestamp + FROM coordination_history + WHERE coordination_request_id = $1 + ORDER BY timestamp ASC + """, + request_id, + ) + + return { + "coordination_request": dict(request), + "history": [dict(h) for h in history], + } + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting coordination request {request_id}: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to get coordination request: {str(e)}" + ) + + +@router.put("/requests/{request_id}/resolve", summary="Resolve coordination request") +async def resolve_coordination_request( + request_id: str = Path(..., description="Coordination request ID"), + resolution: CoordinationResolution = ..., +): + """Resolve a coordination request with a specific plan""" + try: + async with get_db_connection() as conn: + # Check if request exists and is pending + existing_request = await conn.fetchrow( + """ + SELECT id, status FROM coordination_requests WHERE id = $1 + """, + request_id, + ) + + if not existing_request: + raise HTTPException( + status_code=404, detail=f"Coordination request {request_id} not found" + ) + + if existing_request["status"] not in ["pending", "in_progress"]: + raise HTTPException( + status_code=400, + detail=f"Request is already {existing_request['status']}", + ) + + # Update request status and resolution + async with get_db_connection() as conn: + await conn.execute( + """ + UPDATE coordination_requests SET + status = 'resolved', + resolution_plan = $2, + resolved_at = $3, + updated_at = $4 + WHERE id = $1 + """, + request_id, + resolution.resolution_plan, + datetime.utcnow(), + datetime.utcnow(), + ) + + # Log resolution in history + async with get_db_connection() as conn: + await conn.execute( + """ + INSERT INTO coordination_history ( + coordination_request_id, action, actor_id, actor_type, details + ) VALUES ($1, $2, $3, $4, $5) + """, + request_id, + "resolved", + resolution.resolver_id, + "agent", + { + "resolution_plan": resolution.resolution_plan, + "notes": resolution.notes, + }, + ) + + return { + "status": "success", + "request_id": request_id, + "message": "Coordination request resolved successfully", + } + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error resolving coordination request {request_id}: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to resolve coordination request: {str(e)}" + ) + + +# Resource Management Endpoints +@router.post("/resources/allocate", summary="Allocate resources to a product") +async def allocate_resource(allocation: ResourceAllocationCreate): + """Allocate a resource to a specific product""" + try: + async with get_db_connection() as conn: + allocation_id = await conn.fetchval( + """ + INSERT INTO resource_allocations ( + product_id, resource_type, resource_name, allocation_details, + valid_until, status, created_at, updated_at + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + RETURNING id + """, + allocation.product_id, + allocation.resource_type, + allocation.resource_name, + allocation.allocation_details, + allocation.valid_until, + "active", + datetime.utcnow(), + datetime.utcnow(), + ) + + return { + "status": "success", + "allocation_id": str(allocation_id), + "message": "Resource allocated successfully", + } + + except asyncpg.UniqueViolationError: + raise HTTPException( + status_code=409, + detail=f"Resource {allocation.resource_name} of type {allocation.resource_type} is already allocated", + ) + except Exception as e: + logger.error(f"Error allocating resource: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to allocate resource: {str(e)}" + ) + + +@router.get("/resources", summary="List resource allocations") +async def list_resource_allocations( + product_id: Optional[str] = Query(None, description="Filter by product"), + resource_type: Optional[str] = Query(None, description="Filter by resource type"), + status: Optional[str] = Query(None, description="Filter by status"), +): + """List current resource allocations""" + try: + where_conditions = [] + params = [] + param_count = 0 + + if product_id: + param_count += 1 + where_conditions.append(f"product_id = ${param_count}") + params.append(product_id) + + if resource_type: + param_count += 1 + where_conditions.append(f"resource_type = ${param_count}") + params.append(resource_type) + + if status: + param_count += 1 + where_conditions.append(f"status = ${param_count}") + params.append(status) + + where_clause = ( + " WHERE " + " AND ".join(where_conditions) if where_conditions else "" + ) + + query = f""" + SELECT ra.*, pr.name as product_name + FROM resource_allocations ra + LEFT JOIN product_registry pr ON ra.product_id = pr.id + {where_clause} + ORDER BY ra.created_at DESC + """ # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params + + async with get_db_connection() as conn: + allocations = await conn.fetch(query, *params) + + return { + "resource_allocations": [dict(alloc) for alloc in allocations], + "total_count": len(allocations), + } + + except Exception as e: + logger.error(f"Error listing resource allocations: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to list resource allocations: {str(e)}" + ) + + +# Dashboard and Status Endpoints +@router.get("/status", summary="Get overall coordination system status") +async def get_coordination_status(): + """Get overall status of the cross-product coordination system""" + try: + async with get_db_connection() as conn: + # Get request counts by status + request_stats = await conn.fetch(""" + SELECT status, priority, COUNT(*) as count + FROM coordination_requests + GROUP BY status, priority + ORDER BY status, priority + """) + + # Get product count + product_count = await conn.fetchval(""" + SELECT COUNT(*) FROM product_registry + """) + + # Get resource allocation stats + resource_stats = await conn.fetch(""" + SELECT resource_type, status, COUNT(*) as count + FROM resource_allocations + GROUP BY resource_type, status + """) + + # Get recent activity + recent_activity = await conn.fetch(""" + SELECT ch.action, ch.timestamp, cr.title, cr.requesting_product + FROM coordination_history ch + JOIN coordination_requests cr ON ch.coordination_request_id = cr.id + ORDER BY ch.timestamp DESC + LIMIT 10 + """) + + return { + "system_status": "operational", + "registered_products": product_count, + "coordination_request_stats": [dict(stat) for stat in request_stats], + "resource_allocation_stats": [dict(stat) for stat in resource_stats], + "recent_activity": [dict(activity) for activity in recent_activity], + "last_updated": datetime.utcnow().isoformat(), + } + + except Exception as e: + logger.error(f"Error getting coordination status: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to get coordination status: {str(e)}" + ) + + +@router.get("/protocols", summary="List coordination protocols") +async def list_coordination_protocols(): + """List available coordination protocols and procedures""" + try: + async with get_db_connection() as conn: + protocols = await conn.fetch(""" + SELECT protocol_name, coordination_type, scope, procedure_steps, + required_approvals, sla_requirements, is_active, version + FROM coordination_protocols + WHERE is_active = true + ORDER BY protocol_name ASC + """) + + return { + "coordination_protocols": [dict(protocol) for protocol in protocols], + "total_count": len(protocols), + } + + except Exception as e: + logger.error(f"Error listing coordination protocols: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to list coordination protocols: {str(e)}" + ) diff --git a/services/orchestrator/goal_conversation_service.py b/services/orchestrator/goal_conversation_service.py index bf2be4f..ce16acf 100644 --- a/services/orchestrator/goal_conversation_service.py +++ b/services/orchestrator/goal_conversation_service.py @@ -1,1040 +1,1040 @@ -""" -Goal Conversation Management Service for FuzeAgent - -This service manages AI-powered conversations about organizational goals, -enabling collaborative planning, progress reviews, problem-solving, and -strategic adjustments through intelligent dialogue. -""" - -import asyncio -import json -import logging -import uuid -from dataclasses import dataclass -from datetime import datetime, timedelta -from enum import Enum -from typing import Any, Dict, List, Optional, Tuple - -import asyncpg - -logger = logging.getLogger(__name__) - - -class ConversationType(str, Enum): - PLANNING = "planning" - REVIEW = "review" - ADJUSTMENT = "adjustment" - PROBLEM_SOLVING = "problem_solving" - BRAINSTORMING = "brainstorming" - RETROSPECTIVE = "retrospective" - - -class ConversationStatus(str, Enum): - ACTIVE = "active" - ARCHIVED = "archived" - COMPLETED = "completed" - - -class MessageType(str, Enum): - SYSTEM = "system" - AGENT = "agent" - HUMAN = "human" - AI_ANALYSIS = "ai_analysis" - ACTION_ITEM = "action_item" - - -@dataclass -class ConversationMessage: - """Represents a message in a goal conversation""" - - id: str - message_type: MessageType - sender_id: Optional[str] - sender_name: Optional[str] - content: str - metadata: Dict[str, Any] - timestamp: datetime - references: List[str] # Referenced message IDs - reactions: List[Dict[str, Any]] # Message reactions/acknowledgments - - -@dataclass -class ConversationInsight: - """Represents an AI-generated insight from conversation analysis""" - - id: str - insight_type: str # pattern, risk, opportunity, recommendation - title: str - description: str - confidence_score: float - supporting_messages: List[str] - suggested_actions: List[Dict[str, Any]] - generated_at: datetime - - -@dataclass -class ActionItem: - """Represents an action item derived from conversation""" - - id: str - title: str - description: str - assigned_to: Optional[str] - due_date: Optional[datetime] - status: str # pending, in_progress, completed, cancelled - priority: int - source_messages: List[str] - created_at: datetime - completed_at: Optional[datetime] - - -class GoalConversationService: - """ - Manages AI-powered conversations for organizational goal planning, - tracking, and optimization with intelligent insights and action generation. - """ - - def __init__(self, database_url: str): - self.database_url = database_url - self.pool: Optional[asyncpg.Pool] = None - - # Configuration - self.max_conversation_messages = 1000 - self.insight_confidence_threshold = 0.6 - self.auto_action_item_threshold = 0.8 - - # AI conversation templates and prompts - self.conversation_starters = self._initialize_conversation_starters() - self.analysis_prompts = self._initialize_analysis_prompts() - - # Statistics - self.conversations_created = 0 - self.messages_processed = 0 - self.insights_generated = 0 - self.action_items_created = 0 - - async def initialize(self): - """Initialize the goal conversation service""" - logger.info("Initializing GoalConversationService") - - try: - self.pool = await asyncpg.create_pool( - self.database_url, min_size=1, max_size=5, command_timeout=60 - ) - - logger.info("GoalConversationService initialized successfully") - - except Exception as e: - logger.error(f"Failed to initialize GoalConversationService: {e}") - raise - - async def close(self): - """Close database connections""" - if self.pool: - await self.pool.close() - logger.info("GoalConversationService closed") - - async def create_goal_conversation( - self, - goal_id: str, - conversation_type: ConversationType, - conversation_title: str, - initial_context: Optional[Dict[str, Any]] = None, - participants: Optional[List[Dict[str, Any]]] = None, - created_by: Optional[str] = None, - ) -> str: - """Create a new conversation for a goal""" - - conversation_id = str(uuid.uuid4()) - - if initial_context is None: - initial_context = {} - if participants is None: - participants = [] - - try: - async with self.pool.acquire() as conn: - # Get goal context - goal = await conn.fetchrow( - """ - SELECT title, description, goal_type, target_deadline, - progress_percentage, current_value, target_value - FROM organization_goals WHERE id = $1 - """, - goal_id, - ) - - if not goal: - raise ValueError(f"Goal {goal_id} not found") - - # Enhanced context with goal information - enhanced_context = { - **initial_context, - "goal_title": goal["title"], - "goal_type": goal["goal_type"], - "goal_progress": float(goal["progress_percentage"]), - "days_to_deadline": ( - goal["target_deadline"] - datetime.now().date() - ).days, - "conversation_created_at": datetime.now().isoformat(), - } - - # Create conversation - await conn.execute( - """ - INSERT INTO goal_conversations ( - id, goal_id, conversation_type, conversation_title, - conversation_context, participants, status, created_by - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) - """, - conversation_id, - goal_id, - conversation_type.value, - conversation_title, - json.dumps(enhanced_context), - json.dumps(participants), - ConversationStatus.ACTIVE.value, - created_by, - ) - - # Add initial system message with conversation starter - starter_message = self._generate_conversation_starter( - conversation_type, goal, enhanced_context - ) - - await self._add_message( - conversation_id, - MessageType.SYSTEM, - None, - "System", - starter_message, - {"conversation_starter": True}, - ) - - self.conversations_created += 1 - logger.info(f"Created conversation {conversation_id} for goal {goal_id}") - - return conversation_id - - except Exception as e: - logger.error(f"Error creating conversation: {e}") - raise - - async def add_message_to_conversation( - self, - conversation_id: str, - message_type: MessageType, - sender_id: Optional[str], - sender_name: str, - content: str, - metadata: Optional[Dict[str, Any]] = None, - references: Optional[List[str]] = None, - ) -> str: - """Add a message to a conversation""" - - if metadata is None: - metadata = {} - if references is None: - references = [] - - try: - # Add the message - message_id = await self._add_message( - conversation_id, - message_type, - sender_id, - sender_name, - content, - metadata, - references, - ) - - # Trigger conversation analysis for insights - await self._analyze_conversation_for_insights(conversation_id) - - # Update conversation activity timestamp - async with self.pool.acquire() as conn: - await conn.execute( - """ - UPDATE goal_conversations - SET last_activity_at = NOW(), updated_at = NOW() - WHERE id = $1 - """, - conversation_id, - ) - - self.messages_processed += 1 - - return message_id - - except Exception as e: - logger.error(f"Error adding message to conversation {conversation_id}: {e}") - raise - - async def get_conversation(self, conversation_id: str) -> Optional[Dict[str, Any]]: - """Get full conversation with messages, insights, and action items""" - - try: - async with self.pool.acquire() as conn: - # Get conversation details - conversation = await conn.fetchrow( - """ - SELECT gc.*, og.title as goal_title, og.goal_type - FROM goal_conversations gc - JOIN organization_goals og ON gc.goal_id = og.id - WHERE gc.id = $1 - """, - conversation_id, - ) - - if not conversation: - return None - - # Get messages - messages = ( - json.loads(conversation["messages"]) - if conversation["messages"] - else [] - ) - - # Get insights - insights = ( - json.loads(conversation["insights_generated"]) - if conversation["insights_generated"] - else [] - ) - - # Get action items - action_items = ( - json.loads(conversation["action_items"]) - if conversation["action_items"] - else [] - ) - - return { - "id": str(conversation["id"]), - "goal_id": str(conversation["goal_id"]), - "goal_title": conversation["goal_title"], - "conversation_type": conversation["conversation_type"], - "conversation_title": conversation["conversation_title"], - "conversation_summary": conversation["conversation_summary"], - "conversation_context": ( - json.loads(conversation["conversation_context"]) - if conversation["conversation_context"] - else {} - ), - "participants": ( - json.loads(conversation["participants"]) - if conversation["participants"] - else [] - ), - "messages": messages, - "insights_generated": insights, - "action_items": action_items, - "status": conversation["status"], - "last_activity_at": ( - conversation["last_activity_at"].isoformat() - if conversation["last_activity_at"] - else None - ), - "created_at": conversation["created_at"].isoformat(), - "updated_at": conversation["updated_at"].isoformat(), - "message_count": len(messages), - "insight_count": len(insights), - "action_item_count": len(action_items), - } - - except Exception as e: - logger.error(f"Error getting conversation {conversation_id}: {e}") - return None - - async def generate_planning_milestones( - self, conversation_id: str, planning_context: Optional[Dict[str, Any]] = None - ) -> List[Dict[str, Any]]: - """Generate milestone suggestions based on conversation analysis""" - - try: - async with self.pool.acquire() as conn: - # Get conversation and goal context - conversation = await conn.fetchrow( - """ - SELECT gc.*, og.title, og.description, og.goal_type, - og.target_deadline, og.target_value, og.target_unit - FROM goal_conversations gc - JOIN organization_goals og ON gc.goal_id = og.id - WHERE gc.id = $1 - """, - conversation_id, - ) - - if not conversation: - raise ValueError(f"Conversation {conversation_id} not found") - - # Analyze conversation content for milestone ideas - messages = ( - json.loads(conversation["messages"]) - if conversation["messages"] - else [] - ) - milestone_suggestions = self._extract_milestone_ideas_from_conversation( - messages, conversation, planning_context - ) - - # Generate AI-powered milestone recommendations - ai_milestones = await self._generate_ai_milestone_recommendations( - conversation, milestone_suggestions, planning_context - ) - - # Add milestones as insights to the conversation - milestone_insight = { - "id": str(uuid.uuid4()), - "insight_type": "milestone_recommendations", - "title": "AI-Generated Milestone Recommendations", - "description": f"Based on conversation analysis, here are {len(ai_milestones)} recommended milestones", - "confidence_score": 0.85, - "supporting_messages": [ - msg["id"] for msg in messages[-5:] if "id" in msg - ], # Last 5 messages - "suggested_actions": [ - { - "action": "create_milestones", - "description": "Create these milestones for the goal", - "milestones": ai_milestones, - } - ], - "generated_at": datetime.now().isoformat(), - } - - # Update conversation with milestone insight - await self._add_insight_to_conversation( - conversation_id, milestone_insight - ) - - return ai_milestones - - except Exception as e: - logger.error(f"Error generating planning milestones: {e}") - return [] - - async def conduct_progress_review( - self, conversation_id: str, review_period_days: int = 30 - ) -> Dict[str, Any]: - """Conduct AI-powered progress review for a goal conversation""" - - try: - async with self.pool.acquire() as conn: - # Get conversation and goal data - conversation = await conn.fetchrow( - """ - SELECT gc.*, og.* - FROM goal_conversations gc - JOIN organization_goals og ON gc.goal_id = og.id - WHERE gc.id = $1 - """, - conversation_id, - ) - - if not conversation: - raise ValueError(f"Conversation {conversation_id} not found") - - # Get recent progress data - progress_data = await conn.fetch( - """ - SELECT * FROM goal_progress_tracking - WHERE goal_id = $1 - AND recorded_at >= NOW() - INTERVAL '%s days' - ORDER BY recorded_at DESC - """, - str(conversation["goal_id"]), - review_period_days, - ) - - # Get milestones and tasks status - milestone_status = await conn.fetch( - """ - SELECT status, COUNT(*) as count - FROM goal_milestones - WHERE goal_id = $1 - GROUP BY status - """, - str(conversation["goal_id"]), - ) - - task_status = await conn.fetch( - """ - SELECT status, COUNT(*) as count - FROM goal_tasks - WHERE goal_id = $1 - GROUP BY status - """, - str(conversation["goal_id"]), - ) - - # Generate review analysis - review_analysis = { - "review_period_days": review_period_days, - "goal_progress": { - "current_progress": float(conversation["progress_percentage"]), - "target_value": ( - float(conversation["target_value"]) - if conversation["target_value"] - else None - ), - "current_value": ( - float(conversation["current_value"]) - if conversation["current_value"] - else None - ), - "completion_confidence": float( - conversation["completion_confidence"] - ), - }, - "milestone_summary": { - row["status"]: row["count"] for row in milestone_status - }, - "task_summary": { - row["status"]: row["count"] for row in task_status - }, - "progress_trend": self._calculate_progress_trend( - [dict(p) for p in progress_data] - ), - "risk_assessment": self._assess_goal_risks( - conversation, progress_data - ), - "recommendations": self._generate_progress_recommendations( - conversation, progress_data - ), - } - - # Add review as a structured message - review_message = self._format_progress_review_message(review_analysis) - await self._add_message( - conversation_id, - MessageType.AI_ANALYSIS, - None, - "Progress Analyzer", - review_message, - {"review_analysis": review_analysis}, - ) - - return review_analysis - - except Exception as e: - logger.error(f"Error conducting progress review: {e}") - return {"error": str(e)} - - async def extract_action_items_from_conversation( - self, conversation_id: str, auto_assign: bool = True - ) -> List[Dict[str, Any]]: - """Extract and create action items from conversation analysis""" - - try: - async with self.pool.acquire() as conn: - conversation = await conn.fetchrow( - """ - SELECT * FROM goal_conversations WHERE id = $1 - """, - conversation_id, - ) - - if not conversation: - raise ValueError(f"Conversation {conversation_id} not found") - - messages = ( - json.loads(conversation["messages"]) - if conversation["messages"] - else [] - ) - - # Analyze messages for actionable items - potential_actions = self._identify_action_items_in_messages(messages) - - # Convert to action item format - action_items = [] - for action in potential_actions: - if action["confidence"] >= self.auto_action_item_threshold: - action_item = { - "id": str(uuid.uuid4()), - "title": action["title"], - "description": action["description"], - "assigned_to": ( - action.get("assigned_to") if auto_assign else None - ), - "due_date": action.get("due_date"), - "status": "pending", - "priority": action.get("priority", 5), - "source_messages": action["source_messages"], - "created_at": datetime.now().isoformat(), - "confidence_score": action["confidence"], - } - action_items.append(action_item) - - # Update conversation with action items - if action_items: - existing_actions = ( - json.loads(conversation["action_items"]) - if conversation["action_items"] - else [] - ) - all_actions = existing_actions + action_items - - await conn.execute( - """ - UPDATE goal_conversations - SET action_items = $2, updated_at = NOW() - WHERE id = $1 - """, - conversation_id, - json.dumps(all_actions), - ) - - self.action_items_created += len(action_items) - - return action_items - - except Exception as e: - logger.error(f"Error extracting action items: {e}") - return [] - - async def get_goal_conversations( - self, - goal_id: str, - conversation_type: Optional[ConversationType] = None, - status: Optional[ConversationStatus] = None, - limit: int = 10, - ) -> List[Dict[str, Any]]: - """Get conversations for a goal with optional filtering""" - - try: - async with self.pool.acquire() as conn: - where_conditions = ["goal_id = $1"] - params = [goal_id] - param_idx = 2 - - if conversation_type: - where_conditions.append(f"conversation_type = ${param_idx}") - params.append(conversation_type.value) - param_idx += 1 - - if status: - where_conditions.append(f"status = ${param_idx}") - params.append(status.value) - param_idx += 1 - - where_clause = " AND ".join(where_conditions) - - conversations = await conn.fetch( - """ - SELECT id, conversation_type, conversation_title, conversation_summary, - status, last_activity_at, created_at, - COALESCE(array_length(string_to_array(messages::text, '}}'), 1), 0) as message_count, - COALESCE(array_length(string_to_array(action_items::text, '}}'), 1), 0) as action_count - FROM goal_conversations - WHERE {where_clause} - ORDER BY last_activity_at DESC, created_at DESC - LIMIT ${param_idx} - """.format( # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params - where_clause=where_clause, param_idx=param_idx - ), - *params, - limit, - ) - - return [dict(conv) for conv in conversations] - - except Exception as e: - logger.error(f"Error getting conversations for goal {goal_id}: {e}") - return [] - - # Helper methods for conversation management - - async def _add_message( - self, - conversation_id: str, - message_type: MessageType, - sender_id: Optional[str], - sender_name: str, - content: str, - metadata: Dict[str, Any], - references: Optional[List[str]] = None, - ) -> str: - """Add a message to conversation""" - - message_id = str(uuid.uuid4()) - message = { - "id": message_id, - "message_type": message_type.value, - "sender_id": sender_id, - "sender_name": sender_name, - "content": content, - "metadata": metadata, - "timestamp": datetime.now().isoformat(), - "references": references or [], - "reactions": [], - } - - async with self.pool.acquire() as conn: - # Get current messages - current_messages = await conn.fetchval( - """ - SELECT messages FROM goal_conversations WHERE id = $1 - """, - conversation_id, - ) - - messages = json.loads(current_messages) if current_messages else [] - messages.append(message) - - # Limit message history - if len(messages) > self.max_conversation_messages: - messages = messages[-self.max_conversation_messages :] - - # Update conversation - await conn.execute( - """ - UPDATE goal_conversations - SET messages = $2, updated_at = NOW() - WHERE id = $1 - """, - conversation_id, - json.dumps(messages), - ) - - return message_id - - async def _add_insight_to_conversation( - self, conversation_id: str, insight: Dict[str, Any] - ): - """Add an insight to conversation""" - - async with self.pool.acquire() as conn: - # Get current insights - current_insights = await conn.fetchval( - """ - SELECT insights_generated FROM goal_conversations WHERE id = $1 - """, - conversation_id, - ) - - insights = json.loads(current_insights) if current_insights else [] - insights.append(insight) - - # Update conversation - await conn.execute( - """ - UPDATE goal_conversations - SET insights_generated = $2, updated_at = NOW() - WHERE id = $1 - """, - conversation_id, - json.dumps(insights), - ) - - self.insights_generated += 1 - - def _generate_conversation_starter( - self, - conversation_type: ConversationType, - goal: Dict[str, Any], - context: Dict[str, Any], - ) -> str: - """Generate an appropriate conversation starter""" - - starters = self.conversation_starters.get(conversation_type, {}) - - if conversation_type == ConversationType.PLANNING: - return starters["default"].format( - goal_title=goal["title"], - goal_type=goal["goal_type"], - deadline=goal["target_deadline"].strftime("%B %d, %Y"), - target_value=( - goal["target_value"] - if goal["target_value"] - else "defined objectives" - ), - ) - elif conversation_type == ConversationType.REVIEW: - return starters["default"].format( - goal_title=goal["title"], - current_progress=f"{goal['progress_percentage']:.1f}%", - ) - else: - return starters.get("default", f"Let's discuss the {goal['title']} goal.") - - def _initialize_conversation_starters(self) -> Dict[str, Dict[str, str]]: - """Initialize conversation starter templates""" - - return { - ConversationType.PLANNING: { - "default": """Welcome to the strategic planning session for "{goal_title}"! - -🎯 **Goal**: {goal_title} -📊 **Type**: {goal_type} -📅 **Deadline**: {deadline} -🎌 **Target**: {target_value} - -Let's break this goal down into actionable milestones and tasks. Here are some questions to get us started: - -1. **What are the major milestones we need to achieve?** -2. **What dependencies and blockers should we consider?** -3. **Which teams and resources will be involved?** -4. **How should we measure progress along the way?** - -What aspect would you like to focus on first?""" - }, - ConversationType.REVIEW: { - "default": """Time for a progress review of "{goal_title}"! - -📈 **Current Progress**: {current_progress} - -Let's evaluate our progress, identify what's working well, and address any challenges. - -Key areas to discuss: -- Recent achievements and wins -- Current blockers or risks -- Resource allocation and team performance -- Timeline adjustments if needed -- Next steps and priorities - -What would you like to review first?""" - }, - ConversationType.PROBLEM_SOLVING: { - "default": """Problem-solving session for "{goal_title}". - -Let's identify the specific challenges we're facing and work together to find solutions. Please share: -- What specific problems or blockers have emerged? -- What have we tried so far? -- What constraints or requirements should we consider? - -What's the main challenge you'd like to tackle?""" - }, - } - - def _initialize_analysis_prompts(self) -> Dict[str, str]: - """Initialize AI analysis prompts for conversation processing""" - - return { - "extract_milestones": """Analyze this goal conversation and extract potential milestones mentioned or implied. Look for: -- Time-based deliverables or checkpoints -- Measurable objectives or targets -- Dependencies between activities -- Key decision points or reviews - -Return milestones with titles, descriptions, and target dates.""", - "identify_risks": """Analyze this conversation for potential risks, blockers, or concerns mentioned. Look for: -- Resource constraints or availability issues -- Technical challenges or unknowns -- Timeline concerns or dependencies -- Team capacity or skill gaps -- External dependencies or market factors - -Assess the probability and impact of each risk.""", - "extract_actions": """Extract specific action items from this conversation. Look for: -- Tasks or activities that someone needs to do -- Decisions that need to be made -- Information that needs to be gathered -- People who need to be contacted -- Deadlines or time-sensitive items - -Include who should be responsible and when it should be done.""", - } - - # Placeholder implementations for AI analysis methods - - async def _analyze_conversation_for_insights(self, conversation_id: str): - """Analyze conversation for insights (placeholder for AI integration)""" - # This would integrate with an AI service for conversation analysis - pass - - def _extract_milestone_ideas_from_conversation( - self, - messages: List[Dict[str, Any]], - conversation: Dict[str, Any], - planning_context: Optional[Dict[str, Any]], - ) -> List[Dict[str, Any]]: - """Extract milestone ideas from conversation messages""" - # Simplified implementation - would use NLP/AI for real extraction - milestone_keywords = [ - "milestone", - "phase", - "deliverable", - "target", - "deadline", - "complete", - ] - - milestones = [] - for message in messages: - content = message.get("content", "").lower() - if any(keyword in content for keyword in milestone_keywords): - # Extract potential milestone (simplified) - milestones.append( - { - "title": f"Milestone from conversation", - "description": message.get("content", "")[:200], - "source_message": message.get("id"), - "confidence": 0.7, - } - ) - - return milestones[:5] # Return top 5 candidates - - async def _generate_ai_milestone_recommendations( - self, - conversation: Dict[str, Any], - milestone_suggestions: List[Dict[str, Any]], - planning_context: Optional[Dict[str, Any]], - ) -> List[Dict[str, Any]]: - """Generate AI-powered milestone recommendations""" - # Placeholder implementation - would use AI for intelligent milestone generation - - goal_type = conversation["goal_type"] - timeline_days = (conversation["target_deadline"] - datetime.now().date()).days - - # Generate sample milestones based on goal type - if goal_type == "business" and "revenue" in conversation["title"].lower(): - return [ - { - "title": "Foundation Setup", - "description": "Establish initial infrastructure, team structure, and processes", - "target_date": (datetime.now() + timedelta(days=timeline_days // 4)) - .date() - .isoformat(), - "milestone_type": "checkpoint", - }, - { - "title": "Growth Phase Launch", - "description": "Execute marketing campaigns and sales initiatives", - "target_date": (datetime.now() + timedelta(days=timeline_days // 2)) - .date() - .isoformat(), - "milestone_type": "deliverable", - }, - { - "title": "Scale and Optimize", - "description": "Optimize processes and scale operations for target achievement", - "target_date": ( - datetime.now() + timedelta(days=timeline_days * 3 // 4) - ) - .date() - .isoformat(), - "milestone_type": "metric", - }, - ] - - return [] - - # Additional helper methods for conversation analysis - def _calculate_progress_trend(self, progress_data: List[Dict[str, Any]]) -> str: - """Calculate progress trend from historical data""" - if len(progress_data) < 2: - return "insufficient_data" - - # Simple trend calculation - recent_progress = progress_data[0]["progress_percentage"] - older_progress = progress_data[-1]["progress_percentage"] - - if recent_progress > older_progress * 1.1: - return "accelerating" - elif recent_progress < older_progress * 0.9: - return "declining" - else: - return "steady" - - def _assess_goal_risks( - self, goal: Dict[str, Any], progress_data: List[Dict[str, Any]] - ) -> List[Dict[str, Any]]: - """Assess risks based on goal and progress data""" - risks = [] - - # Timeline risk - days_remaining = (goal["target_deadline"] - datetime.now().date()).days - progress = float(goal["progress_percentage"]) - - if days_remaining < 30 and progress < 70: - risks.append( - { - "type": "timeline_risk", - "severity": "high", - "description": "Goal progress is behind schedule with limited time remaining", - } - ) - - return risks - - def _generate_progress_recommendations( - self, goal: Dict[str, Any], progress_data: List[Dict[str, Any]] - ) -> List[str]: - """Generate recommendations based on progress analysis""" - recommendations = [] - - progress = float(goal["progress_percentage"]) - - if progress < 25: - recommendations.append( - "Consider breaking down remaining work into smaller, more manageable tasks" - ) - - if progress > 75: - recommendations.append( - "Focus on final quality checks and prepare for goal completion" - ) - - return recommendations - - def _identify_action_items_in_messages( - self, messages: List[Dict[str, Any]] - ) -> List[Dict[str, Any]]: - """Identify potential action items from messages""" - # Simplified implementation - would use NLP for real extraction - action_keywords = [ - "need to", - "should", - "must", - "will", - "action", - "task", - "todo", - ] - - actions = [] - for message in messages: - content = message.get("content", "").lower() - if any(keyword in content for keyword in action_keywords): - actions.append( - { - "title": f"Action item from conversation", - "description": message.get("content", "")[:200], - "source_messages": [message.get("id")], - "confidence": 0.8, - "priority": 5, - } - ) - - return actions[:10] # Return top 10 candidates - - def _format_progress_review_message(self, review_analysis: Dict[str, Any]) -> str: - """Format progress review analysis as a readable message""" - - progress = review_analysis["goal_progress"]["current_progress"] - trend = review_analysis["progress_trend"] - - message = f"""## Progress Review Summary - -**Current Progress**: {progress:.1f}% -**Trend**: {trend.replace('_', ' ').title()} - -### Key Metrics -""" - - if review_analysis["milestone_summary"]: - message += "\n**Milestones:**\n" - for status, count in review_analysis["milestone_summary"].items(): - message += f"- {status.replace('_', ' ').title()}: {count}\n" - - if review_analysis["recommendations"]: - message += "\n### Recommendations\n" - for i, rec in enumerate(review_analysis["recommendations"], 1): - message += f"{i}. {rec}\n" - - return message +""" +Goal Conversation Management Service for FuzeAgent + +This service manages AI-powered conversations about organizational goals, +enabling collaborative planning, progress reviews, problem-solving, and +strategic adjustments through intelligent dialogue. +""" + +import asyncio +import json +import logging +import uuid +from dataclasses import dataclass +from datetime import datetime, timedelta +from enum import Enum +from typing import Any, Dict, List, Optional, Tuple + +import asyncpg + +logger = logging.getLogger(__name__) + + +class ConversationType(str, Enum): + PLANNING = "planning" + REVIEW = "review" + ADJUSTMENT = "adjustment" + PROBLEM_SOLVING = "problem_solving" + BRAINSTORMING = "brainstorming" + RETROSPECTIVE = "retrospective" + + +class ConversationStatus(str, Enum): + ACTIVE = "active" + ARCHIVED = "archived" + COMPLETED = "completed" + + +class MessageType(str, Enum): + SYSTEM = "system" + AGENT = "agent" + HUMAN = "human" + AI_ANALYSIS = "ai_analysis" + ACTION_ITEM = "action_item" + + +@dataclass +class ConversationMessage: + """Represents a message in a goal conversation""" + + id: str + message_type: MessageType + sender_id: Optional[str] + sender_name: Optional[str] + content: str + metadata: Dict[str, Any] + timestamp: datetime + references: List[str] # Referenced message IDs + reactions: List[Dict[str, Any]] # Message reactions/acknowledgments + + +@dataclass +class ConversationInsight: + """Represents an AI-generated insight from conversation analysis""" + + id: str + insight_type: str # pattern, risk, opportunity, recommendation + title: str + description: str + confidence_score: float + supporting_messages: List[str] + suggested_actions: List[Dict[str, Any]] + generated_at: datetime + + +@dataclass +class ActionItem: + """Represents an action item derived from conversation""" + + id: str + title: str + description: str + assigned_to: Optional[str] + due_date: Optional[datetime] + status: str # pending, in_progress, completed, cancelled + priority: int + source_messages: List[str] + created_at: datetime + completed_at: Optional[datetime] + + +class GoalConversationService: + """ + Manages AI-powered conversations for organizational goal planning, + tracking, and optimization with intelligent insights and action generation. + """ + + def __init__(self, database_url: str): + self.database_url = database_url + self.pool: Optional[asyncpg.Pool] = None + + # Configuration + self.max_conversation_messages = 1000 + self.insight_confidence_threshold = 0.6 + self.auto_action_item_threshold = 0.8 + + # AI conversation templates and prompts + self.conversation_starters = self._initialize_conversation_starters() + self.analysis_prompts = self._initialize_analysis_prompts() + + # Statistics + self.conversations_created = 0 + self.messages_processed = 0 + self.insights_generated = 0 + self.action_items_created = 0 + + async def initialize(self): + """Initialize the goal conversation service""" + logger.info("Initializing GoalConversationService") + + try: + self.pool = await asyncpg.create_pool( + self.database_url, min_size=1, max_size=5, command_timeout=60 + ) + + logger.info("GoalConversationService initialized successfully") + + except Exception as e: + logger.error(f"Failed to initialize GoalConversationService: {e}") + raise + + async def close(self): + """Close database connections""" + if self.pool: + await self.pool.close() + logger.info("GoalConversationService closed") + + async def create_goal_conversation( + self, + goal_id: str, + conversation_type: ConversationType, + conversation_title: str, + initial_context: Optional[Dict[str, Any]] = None, + participants: Optional[List[Dict[str, Any]]] = None, + created_by: Optional[str] = None, + ) -> str: + """Create a new conversation for a goal""" + + conversation_id = str(uuid.uuid4()) + + if initial_context is None: + initial_context = {} + if participants is None: + participants = [] + + try: + async with self.pool.acquire() as conn: + # Get goal context + goal = await conn.fetchrow( + """ + SELECT title, description, goal_type, target_deadline, + progress_percentage, current_value, target_value + FROM organization_goals WHERE id = $1 + """, + goal_id, + ) + + if not goal: + raise ValueError(f"Goal {goal_id} not found") + + # Enhanced context with goal information + enhanced_context = { + **initial_context, + "goal_title": goal["title"], + "goal_type": goal["goal_type"], + "goal_progress": float(goal["progress_percentage"]), + "days_to_deadline": ( + goal["target_deadline"] - datetime.now().date() + ).days, + "conversation_created_at": datetime.now().isoformat(), + } + + # Create conversation + await conn.execute( + """ + INSERT INTO goal_conversations ( + id, goal_id, conversation_type, conversation_title, + conversation_context, participants, status, created_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + """, + conversation_id, + goal_id, + conversation_type.value, + conversation_title, + json.dumps(enhanced_context), + json.dumps(participants), + ConversationStatus.ACTIVE.value, + created_by, + ) + + # Add initial system message with conversation starter + starter_message = self._generate_conversation_starter( + conversation_type, goal, enhanced_context + ) + + await self._add_message( + conversation_id, + MessageType.SYSTEM, + None, + "System", + starter_message, + {"conversation_starter": True}, + ) + + self.conversations_created += 1 + logger.info(f"Created conversation {conversation_id} for goal {goal_id}") + + return conversation_id + + except Exception as e: + logger.error(f"Error creating conversation: {e}") + raise + + async def add_message_to_conversation( + self, + conversation_id: str, + message_type: MessageType, + sender_id: Optional[str], + sender_name: str, + content: str, + metadata: Optional[Dict[str, Any]] = None, + references: Optional[List[str]] = None, + ) -> str: + """Add a message to a conversation""" + + if metadata is None: + metadata = {} + if references is None: + references = [] + + try: + # Add the message + message_id = await self._add_message( + conversation_id, + message_type, + sender_id, + sender_name, + content, + metadata, + references, + ) + + # Trigger conversation analysis for insights + await self._analyze_conversation_for_insights(conversation_id) + + # Update conversation activity timestamp + async with self.pool.acquire() as conn: + await conn.execute( + """ + UPDATE goal_conversations + SET last_activity_at = NOW(), updated_at = NOW() + WHERE id = $1 + """, + conversation_id, + ) + + self.messages_processed += 1 + + return message_id + + except Exception as e: + logger.error(f"Error adding message to conversation {conversation_id}: {e}") + raise + + async def get_conversation(self, conversation_id: str) -> Optional[Dict[str, Any]]: + """Get full conversation with messages, insights, and action items""" + + try: + async with self.pool.acquire() as conn: + # Get conversation details + conversation = await conn.fetchrow( + """ + SELECT gc.*, og.title as goal_title, og.goal_type + FROM goal_conversations gc + JOIN organization_goals og ON gc.goal_id = og.id + WHERE gc.id = $1 + """, + conversation_id, + ) + + if not conversation: + return None + + # Get messages + messages = ( + json.loads(conversation["messages"]) + if conversation["messages"] + else [] + ) + + # Get insights + insights = ( + json.loads(conversation["insights_generated"]) + if conversation["insights_generated"] + else [] + ) + + # Get action items + action_items = ( + json.loads(conversation["action_items"]) + if conversation["action_items"] + else [] + ) + + return { + "id": str(conversation["id"]), + "goal_id": str(conversation["goal_id"]), + "goal_title": conversation["goal_title"], + "conversation_type": conversation["conversation_type"], + "conversation_title": conversation["conversation_title"], + "conversation_summary": conversation["conversation_summary"], + "conversation_context": ( + json.loads(conversation["conversation_context"]) + if conversation["conversation_context"] + else {} + ), + "participants": ( + json.loads(conversation["participants"]) + if conversation["participants"] + else [] + ), + "messages": messages, + "insights_generated": insights, + "action_items": action_items, + "status": conversation["status"], + "last_activity_at": ( + conversation["last_activity_at"].isoformat() + if conversation["last_activity_at"] + else None + ), + "created_at": conversation["created_at"].isoformat(), + "updated_at": conversation["updated_at"].isoformat(), + "message_count": len(messages), + "insight_count": len(insights), + "action_item_count": len(action_items), + } + + except Exception as e: + logger.error(f"Error getting conversation {conversation_id}: {e}") + return None + + async def generate_planning_milestones( + self, conversation_id: str, planning_context: Optional[Dict[str, Any]] = None + ) -> List[Dict[str, Any]]: + """Generate milestone suggestions based on conversation analysis""" + + try: + async with self.pool.acquire() as conn: + # Get conversation and goal context + conversation = await conn.fetchrow( + """ + SELECT gc.*, og.title, og.description, og.goal_type, + og.target_deadline, og.target_value, og.target_unit + FROM goal_conversations gc + JOIN organization_goals og ON gc.goal_id = og.id + WHERE gc.id = $1 + """, + conversation_id, + ) + + if not conversation: + raise ValueError(f"Conversation {conversation_id} not found") + + # Analyze conversation content for milestone ideas + messages = ( + json.loads(conversation["messages"]) + if conversation["messages"] + else [] + ) + milestone_suggestions = self._extract_milestone_ideas_from_conversation( + messages, conversation, planning_context + ) + + # Generate AI-powered milestone recommendations + ai_milestones = await self._generate_ai_milestone_recommendations( + conversation, milestone_suggestions, planning_context + ) + + # Add milestones as insights to the conversation + milestone_insight = { + "id": str(uuid.uuid4()), + "insight_type": "milestone_recommendations", + "title": "AI-Generated Milestone Recommendations", + "description": f"Based on conversation analysis, here are {len(ai_milestones)} recommended milestones", + "confidence_score": 0.85, + "supporting_messages": [ + msg["id"] for msg in messages[-5:] if "id" in msg + ], # Last 5 messages + "suggested_actions": [ + { + "action": "create_milestones", + "description": "Create these milestones for the goal", + "milestones": ai_milestones, + } + ], + "generated_at": datetime.now().isoformat(), + } + + # Update conversation with milestone insight + await self._add_insight_to_conversation( + conversation_id, milestone_insight + ) + + return ai_milestones + + except Exception as e: + logger.error(f"Error generating planning milestones: {e}") + return [] + + async def conduct_progress_review( + self, conversation_id: str, review_period_days: int = 30 + ) -> Dict[str, Any]: + """Conduct AI-powered progress review for a goal conversation""" + + try: + async with self.pool.acquire() as conn: + # Get conversation and goal data + conversation = await conn.fetchrow( + """ + SELECT gc.*, og.* + FROM goal_conversations gc + JOIN organization_goals og ON gc.goal_id = og.id + WHERE gc.id = $1 + """, + conversation_id, + ) + + if not conversation: + raise ValueError(f"Conversation {conversation_id} not found") + + # Get recent progress data + progress_data = await conn.fetch( + """ + SELECT * FROM goal_progress_tracking + WHERE goal_id = $1 + AND recorded_at >= NOW() - INTERVAL '%s days' + ORDER BY recorded_at DESC + """, + str(conversation["goal_id"]), + review_period_days, + ) + + # Get milestones and tasks status + milestone_status = await conn.fetch( + """ + SELECT status, COUNT(*) as count + FROM goal_milestones + WHERE goal_id = $1 + GROUP BY status + """, + str(conversation["goal_id"]), + ) + + task_status = await conn.fetch( + """ + SELECT status, COUNT(*) as count + FROM goal_tasks + WHERE goal_id = $1 + GROUP BY status + """, + str(conversation["goal_id"]), + ) + + # Generate review analysis + review_analysis = { + "review_period_days": review_period_days, + "goal_progress": { + "current_progress": float(conversation["progress_percentage"]), + "target_value": ( + float(conversation["target_value"]) + if conversation["target_value"] + else None + ), + "current_value": ( + float(conversation["current_value"]) + if conversation["current_value"] + else None + ), + "completion_confidence": float( + conversation["completion_confidence"] + ), + }, + "milestone_summary": { + row["status"]: row["count"] for row in milestone_status + }, + "task_summary": { + row["status"]: row["count"] for row in task_status + }, + "progress_trend": self._calculate_progress_trend( + [dict(p) for p in progress_data] + ), + "risk_assessment": self._assess_goal_risks( + conversation, progress_data + ), + "recommendations": self._generate_progress_recommendations( + conversation, progress_data + ), + } + + # Add review as a structured message + review_message = self._format_progress_review_message(review_analysis) + await self._add_message( + conversation_id, + MessageType.AI_ANALYSIS, + None, + "Progress Analyzer", + review_message, + {"review_analysis": review_analysis}, + ) + + return review_analysis + + except Exception as e: + logger.error(f"Error conducting progress review: {e}") + return {"error": str(e)} + + async def extract_action_items_from_conversation( + self, conversation_id: str, auto_assign: bool = True + ) -> List[Dict[str, Any]]: + """Extract and create action items from conversation analysis""" + + try: + async with self.pool.acquire() as conn: + conversation = await conn.fetchrow( + """ + SELECT * FROM goal_conversations WHERE id = $1 + """, + conversation_id, + ) + + if not conversation: + raise ValueError(f"Conversation {conversation_id} not found") + + messages = ( + json.loads(conversation["messages"]) + if conversation["messages"] + else [] + ) + + # Analyze messages for actionable items + potential_actions = self._identify_action_items_in_messages(messages) + + # Convert to action item format + action_items = [] + for action in potential_actions: + if action["confidence"] >= self.auto_action_item_threshold: + action_item = { + "id": str(uuid.uuid4()), + "title": action["title"], + "description": action["description"], + "assigned_to": ( + action.get("assigned_to") if auto_assign else None + ), + "due_date": action.get("due_date"), + "status": "pending", + "priority": action.get("priority", 5), + "source_messages": action["source_messages"], + "created_at": datetime.now().isoformat(), + "confidence_score": action["confidence"], + } + action_items.append(action_item) + + # Update conversation with action items + if action_items: + existing_actions = ( + json.loads(conversation["action_items"]) + if conversation["action_items"] + else [] + ) + all_actions = existing_actions + action_items + + await conn.execute( + """ + UPDATE goal_conversations + SET action_items = $2, updated_at = NOW() + WHERE id = $1 + """, + conversation_id, + json.dumps(all_actions), + ) + + self.action_items_created += len(action_items) + + return action_items + + except Exception as e: + logger.error(f"Error extracting action items: {e}") + return [] + + async def get_goal_conversations( + self, + goal_id: str, + conversation_type: Optional[ConversationType] = None, + status: Optional[ConversationStatus] = None, + limit: int = 10, + ) -> List[Dict[str, Any]]: + """Get conversations for a goal with optional filtering""" + + try: + async with self.pool.acquire() as conn: + where_conditions = ["goal_id = $1"] + params = [goal_id] + param_idx = 2 + + if conversation_type: + where_conditions.append(f"conversation_type = ${param_idx}") + params.append(conversation_type.value) + param_idx += 1 + + if status: + where_conditions.append(f"status = ${param_idx}") + params.append(status.value) + param_idx += 1 + + where_clause = " AND ".join(where_conditions) + + conversations = await conn.fetch( + """ + SELECT id, conversation_type, conversation_title, conversation_summary, + status, last_activity_at, created_at, + COALESCE(array_length(string_to_array(messages::text, '}}'), 1), 0) as message_count, + COALESCE(array_length(string_to_array(action_items::text, '}}'), 1), 0) as action_count + FROM goal_conversations + WHERE {where_clause} + ORDER BY last_activity_at DESC, created_at DESC + LIMIT ${param_idx} + """.format( # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params + where_clause=where_clause, param_idx=param_idx + ), + *params, + limit, + ) + + return [dict(conv) for conv in conversations] + + except Exception as e: + logger.error(f"Error getting conversations for goal {goal_id}: {e}") + return [] + + # Helper methods for conversation management + + async def _add_message( + self, + conversation_id: str, + message_type: MessageType, + sender_id: Optional[str], + sender_name: str, + content: str, + metadata: Dict[str, Any], + references: Optional[List[str]] = None, + ) -> str: + """Add a message to conversation""" + + message_id = str(uuid.uuid4()) + message = { + "id": message_id, + "message_type": message_type.value, + "sender_id": sender_id, + "sender_name": sender_name, + "content": content, + "metadata": metadata, + "timestamp": datetime.now().isoformat(), + "references": references or [], + "reactions": [], + } + + async with self.pool.acquire() as conn: + # Get current messages + current_messages = await conn.fetchval( + """ + SELECT messages FROM goal_conversations WHERE id = $1 + """, + conversation_id, + ) + + messages = json.loads(current_messages) if current_messages else [] + messages.append(message) + + # Limit message history + if len(messages) > self.max_conversation_messages: + messages = messages[-self.max_conversation_messages :] + + # Update conversation + await conn.execute( + """ + UPDATE goal_conversations + SET messages = $2, updated_at = NOW() + WHERE id = $1 + """, + conversation_id, + json.dumps(messages), + ) + + return message_id + + async def _add_insight_to_conversation( + self, conversation_id: str, insight: Dict[str, Any] + ): + """Add an insight to conversation""" + + async with self.pool.acquire() as conn: + # Get current insights + current_insights = await conn.fetchval( + """ + SELECT insights_generated FROM goal_conversations WHERE id = $1 + """, + conversation_id, + ) + + insights = json.loads(current_insights) if current_insights else [] + insights.append(insight) + + # Update conversation + await conn.execute( + """ + UPDATE goal_conversations + SET insights_generated = $2, updated_at = NOW() + WHERE id = $1 + """, + conversation_id, + json.dumps(insights), + ) + + self.insights_generated += 1 + + def _generate_conversation_starter( + self, + conversation_type: ConversationType, + goal: Dict[str, Any], + context: Dict[str, Any], + ) -> str: + """Generate an appropriate conversation starter""" + + starters = self.conversation_starters.get(conversation_type, {}) + + if conversation_type == ConversationType.PLANNING: + return starters["default"].format( + goal_title=goal["title"], + goal_type=goal["goal_type"], + deadline=goal["target_deadline"].strftime("%B %d, %Y"), + target_value=( + goal["target_value"] + if goal["target_value"] + else "defined objectives" + ), + ) + elif conversation_type == ConversationType.REVIEW: + return starters["default"].format( + goal_title=goal["title"], + current_progress=f"{goal['progress_percentage']:.1f}%", + ) + else: + return starters.get("default", f"Let's discuss the {goal['title']} goal.") + + def _initialize_conversation_starters(self) -> Dict[str, Dict[str, str]]: + """Initialize conversation starter templates""" + + return { + ConversationType.PLANNING: { + "default": """Welcome to the strategic planning session for "{goal_title}"! + +🎯 **Goal**: {goal_title} +📊 **Type**: {goal_type} +📅 **Deadline**: {deadline} +🎌 **Target**: {target_value} + +Let's break this goal down into actionable milestones and tasks. Here are some questions to get us started: + +1. **What are the major milestones we need to achieve?** +2. **What dependencies and blockers should we consider?** +3. **Which teams and resources will be involved?** +4. **How should we measure progress along the way?** + +What aspect would you like to focus on first?""" + }, + ConversationType.REVIEW: { + "default": """Time for a progress review of "{goal_title}"! + +📈 **Current Progress**: {current_progress} + +Let's evaluate our progress, identify what's working well, and address any challenges. + +Key areas to discuss: +- Recent achievements and wins +- Current blockers or risks +- Resource allocation and team performance +- Timeline adjustments if needed +- Next steps and priorities + +What would you like to review first?""" + }, + ConversationType.PROBLEM_SOLVING: { + "default": """Problem-solving session for "{goal_title}". + +Let's identify the specific challenges we're facing and work together to find solutions. Please share: +- What specific problems or blockers have emerged? +- What have we tried so far? +- What constraints or requirements should we consider? + +What's the main challenge you'd like to tackle?""" + }, + } + + def _initialize_analysis_prompts(self) -> Dict[str, str]: + """Initialize AI analysis prompts for conversation processing""" + + return { + "extract_milestones": """Analyze this goal conversation and extract potential milestones mentioned or implied. Look for: +- Time-based deliverables or checkpoints +- Measurable objectives or targets +- Dependencies between activities +- Key decision points or reviews + +Return milestones with titles, descriptions, and target dates.""", + "identify_risks": """Analyze this conversation for potential risks, blockers, or concerns mentioned. Look for: +- Resource constraints or availability issues +- Technical challenges or unknowns +- Timeline concerns or dependencies +- Team capacity or skill gaps +- External dependencies or market factors + +Assess the probability and impact of each risk.""", + "extract_actions": """Extract specific action items from this conversation. Look for: +- Tasks or activities that someone needs to do +- Decisions that need to be made +- Information that needs to be gathered +- People who need to be contacted +- Deadlines or time-sensitive items + +Include who should be responsible and when it should be done.""", + } + + # Placeholder implementations for AI analysis methods + + async def _analyze_conversation_for_insights(self, conversation_id: str): + """Analyze conversation for insights (placeholder for AI integration)""" + # This would integrate with an AI service for conversation analysis + pass + + def _extract_milestone_ideas_from_conversation( + self, + messages: List[Dict[str, Any]], + conversation: Dict[str, Any], + planning_context: Optional[Dict[str, Any]], + ) -> List[Dict[str, Any]]: + """Extract milestone ideas from conversation messages""" + # Simplified implementation - would use NLP/AI for real extraction + milestone_keywords = [ + "milestone", + "phase", + "deliverable", + "target", + "deadline", + "complete", + ] + + milestones = [] + for message in messages: + content = message.get("content", "").lower() + if any(keyword in content for keyword in milestone_keywords): + # Extract potential milestone (simplified) + milestones.append( + { + "title": f"Milestone from conversation", + "description": message.get("content", "")[:200], + "source_message": message.get("id"), + "confidence": 0.7, + } + ) + + return milestones[:5] # Return top 5 candidates + + async def _generate_ai_milestone_recommendations( + self, + conversation: Dict[str, Any], + milestone_suggestions: List[Dict[str, Any]], + planning_context: Optional[Dict[str, Any]], + ) -> List[Dict[str, Any]]: + """Generate AI-powered milestone recommendations""" + # Placeholder implementation - would use AI for intelligent milestone generation + + goal_type = conversation["goal_type"] + timeline_days = (conversation["target_deadline"] - datetime.now().date()).days + + # Generate sample milestones based on goal type + if goal_type == "business" and "revenue" in conversation["title"].lower(): + return [ + { + "title": "Foundation Setup", + "description": "Establish initial infrastructure, team structure, and processes", + "target_date": (datetime.now() + timedelta(days=timeline_days // 4)) + .date() + .isoformat(), + "milestone_type": "checkpoint", + }, + { + "title": "Growth Phase Launch", + "description": "Execute marketing campaigns and sales initiatives", + "target_date": (datetime.now() + timedelta(days=timeline_days // 2)) + .date() + .isoformat(), + "milestone_type": "deliverable", + }, + { + "title": "Scale and Optimize", + "description": "Optimize processes and scale operations for target achievement", + "target_date": ( + datetime.now() + timedelta(days=timeline_days * 3 // 4) + ) + .date() + .isoformat(), + "milestone_type": "metric", + }, + ] + + return [] + + # Additional helper methods for conversation analysis + def _calculate_progress_trend(self, progress_data: List[Dict[str, Any]]) -> str: + """Calculate progress trend from historical data""" + if len(progress_data) < 2: + return "insufficient_data" + + # Simple trend calculation + recent_progress = progress_data[0]["progress_percentage"] + older_progress = progress_data[-1]["progress_percentage"] + + if recent_progress > older_progress * 1.1: + return "accelerating" + elif recent_progress < older_progress * 0.9: + return "declining" + else: + return "steady" + + def _assess_goal_risks( + self, goal: Dict[str, Any], progress_data: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + """Assess risks based on goal and progress data""" + risks = [] + + # Timeline risk + days_remaining = (goal["target_deadline"] - datetime.now().date()).days + progress = float(goal["progress_percentage"]) + + if days_remaining < 30 and progress < 70: + risks.append( + { + "type": "timeline_risk", + "severity": "high", + "description": "Goal progress is behind schedule with limited time remaining", + } + ) + + return risks + + def _generate_progress_recommendations( + self, goal: Dict[str, Any], progress_data: List[Dict[str, Any]] + ) -> List[str]: + """Generate recommendations based on progress analysis""" + recommendations = [] + + progress = float(goal["progress_percentage"]) + + if progress < 25: + recommendations.append( + "Consider breaking down remaining work into smaller, more manageable tasks" + ) + + if progress > 75: + recommendations.append( + "Focus on final quality checks and prepare for goal completion" + ) + + return recommendations + + def _identify_action_items_in_messages( + self, messages: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + """Identify potential action items from messages""" + # Simplified implementation - would use NLP for real extraction + action_keywords = [ + "need to", + "should", + "must", + "will", + "action", + "task", + "todo", + ] + + actions = [] + for message in messages: + content = message.get("content", "").lower() + if any(keyword in content for keyword in action_keywords): + actions.append( + { + "title": f"Action item from conversation", + "description": message.get("content", "")[:200], + "source_messages": [message.get("id")], + "confidence": 0.8, + "priority": 5, + } + ) + + return actions[:10] # Return top 10 candidates + + def _format_progress_review_message(self, review_analysis: Dict[str, Any]) -> str: + """Format progress review analysis as a readable message""" + + progress = review_analysis["goal_progress"]["current_progress"] + trend = review_analysis["progress_trend"] + + message = f"""## Progress Review Summary + +**Current Progress**: {progress:.1f}% +**Trend**: {trend.replace('_', ' ').title()} + +### Key Metrics +""" + + if review_analysis["milestone_summary"]: + message += "\n**Milestones:**\n" + for status, count in review_analysis["milestone_summary"].items(): + message += f"- {status.replace('_', ' ').title()}: {count}\n" + + if review_analysis["recommendations"]: + message += "\n### Recommendations\n" + for i, rec in enumerate(review_analysis["recommendations"], 1): + message += f"{i}. {rec}\n" + + return message diff --git a/services/orchestrator/hierarchy_endpoints.py b/services/orchestrator/hierarchy_endpoints.py index 083d97f..95d2d86 100644 --- a/services/orchestrator/hierarchy_endpoints.py +++ b/services/orchestrator/hierarchy_endpoints.py @@ -1,393 +1,393 @@ -import asyncio -import json -from typing import Any, Dict, List, Optional - -import httpx -from fastapi import APIRouter, HTTPException - -from database import DatabaseManager - -router = APIRouter(prefix="/hierarchy", tags=["hierarchy"]) - - -@router.get("/visualization") -async def get_hierarchy_visualization(): - """ - Get complete organizational hierarchy for visualization with ReactFlow/GoJS - - Returns structured data optimized for hierarchical visualization libraries: - - Organizations as root nodes - - Teams as intermediate nodes - - Agents as leaf nodes - - Relationships and positioning data - """ - try: - # Get data from hierarchy API - async with httpx.AsyncClient() as client: - # Get all organizations - orgs_response = await client.get("http://localhost:8006/organizations") - organizations = orgs_response.json() - - # Get all teams - teams_response = await client.get("http://localhost:8006/teams") - all_teams = teams_response.json() - - # Get all agents - agents_response = await client.get("http://localhost:8006/agents") - all_agents = agents_response.json() - - if not organizations: - return {"nodes": [], "edges": [], "message": "No organizations found"} - - nodes = [] - edges = [] - y_offset = 0 - - for org_idx, org in enumerate(organizations): - org_id = org["id"] - org_node_id = f"org-{org_id}" - - # Add organization node - nodes.append( - { - "id": org_node_id, - "type": "organization", - "data": { - "label": org["name"], - "description": org.get("description", ""), - "type": "Organization", - "settings": org.get("settings", {}), - "entity_id": org_id, - }, - "position": {"x": org_idx * 800, "y": y_offset}, - "style": { - "background": "#1e40af", - "color": "white", - "border": "2px solid #1e3a8a", - "borderRadius": "12px", - "padding": "12px", - "minWidth": "200px", - }, - } - ) - - # Filter teams for this organization - teams = [t for t in all_teams if t.get("organization_id") == org_id] - team_y_offset = y_offset + 150 - - for team_idx, team in enumerate(teams): - team_id = team["id"] - team_node_id = f"team-{team_id}" - - # Add team node - nodes.append( - { - "id": team_node_id, - "type": "team", - "data": { - "label": team["name"], - "description": team.get("description", ""), - "type": f"Team ({team.get('team_type', 'general')})", - "settings": team.get("settings", {}), - "entity_id": team_id, - "organization_id": org_id, - }, - "position": { - "x": org_idx * 800 + (team_idx % 3) * 250 - 250, - "y": team_y_offset + (team_idx // 3) * 120, - }, - "style": { - "background": "#059669", - "color": "white", - "border": "2px solid #047857", - "borderRadius": "8px", - "padding": "10px", - "minWidth": "180px", - }, - } - ) - - # Add edge from organization to team - edges.append( - { - "id": f"edge-{org_node_id}-{team_node_id}", - "source": org_node_id, - "target": team_node_id, - "type": "smoothstep", - "style": {"stroke": "#64748b", "strokeWidth": 2}, - "markerEnd": {"type": "arrowclosed", "color": "#64748b"}, - } - ) - - # Filter agents for this team - note: need to check team_id in agents - agents = [ - a - for a in all_agents - if a.get("team_id") == team_id - or (hasattr(a, "config") and a.config.get("team_id") == team_id) - ] - # Fallback: if no team_id in agent data, we'll have limited agents - if not agents and all_agents: - # For now, include some agents if they don't have team assignments - agents = all_agents[:2] # Include first 2 agents as examples - - agent_y_offset = team_y_offset + (team_idx // 3) * 120 + 100 - - for agent_idx, agent in enumerate(agents): - agent_id = agent["id"] - agent_node_id = f"agent-{agent_id}" - - # Determine agent color by type - agent_colors = { - "executive": {"bg": "#dc2626", "border": "#b91c1c"}, - "developer": {"bg": "#2563eb", "border": "#1d4ed8"}, - "marketing": {"bg": "#7c3aed", "border": "#6d28d9"}, - "sales": {"bg": "#ea580c", "border": "#c2410c"}, - "qa": {"bg": "#16a34a", "border": "#15803d"}, - "devops": {"bg": "#0891b2", "border": "#0e7490"}, - "designer": {"bg": "#e11d48", "border": "#be185d"}, - } - - agent_type = agent.get("type", "developer") - colors = agent_colors.get( - agent_type, {"bg": "#6b7280", "border": "#4b5563"} - ) - - # Add agent node - nodes.append( - { - "id": agent_node_id, - "type": "agent", - "data": { - "label": agent["name"], - "role": agent.get("role", ""), - "type": f"Agent ({agent_type})", - "status": agent.get("status", "active"), - "config": agent.get("config", {}), - "entity_id": agent_id, - "team_id": team_id, - "organization_id": org_id, - }, - "position": { - "x": org_idx * 800 - + (team_idx % 3) * 250 - - 250 - + (agent_idx % 2) * 120 - - 60, - "y": agent_y_offset + (agent_idx // 2) * 80, - }, - "style": { - "background": colors["bg"], - "color": "white", - "border": f"2px solid {colors['border']}", - "borderRadius": "6px", - "padding": "8px", - "minWidth": "150px", - }, - } - ) - - # Add edge from team to agent - edges.append( - { - "id": f"edge-{team_node_id}-{agent_node_id}", - "source": team_node_id, - "target": agent_node_id, - "type": "smoothstep", - "style": {"stroke": "#94a3b8", "strokeWidth": 1.5}, - "markerEnd": {"type": "arrowclosed", "color": "#94a3b8"}, - } - ) - - return { - "nodes": nodes, - "edges": edges, - "metadata": { - "total_organizations": len(organizations), - "total_teams": len(all_teams), - "total_agents": len(all_agents), - "generated_at": "2025-08-06T10:00:00Z", - }, - } - - except Exception as e: - raise HTTPException( - status_code=500, - detail=f"Failed to generate hierarchy visualization: {str(e)}", - ) - - -@router.get("/stats") -async def get_hierarchy_stats(): - """Get comprehensive hierarchy statistics""" - try: - organizations = await DatabaseManager.get_organizations() - - if not organizations: - return {"organizations": 0, "teams": 0, "agents": 0, "by_organization": []} - - stats = { - "organizations": len(organizations), - "teams": 0, - "agents": 0, - "by_organization": [], - "agent_types": {}, - "team_types": {}, - } - - for org in organizations: - org_id = org["id"] - teams = await DatabaseManager.get_teams(org_id) - org_agent_count = 0 - org_teams_by_type = {} - org_agents_by_type = {} - - for team in teams: - team_id = team["id"] - team_type = team.get("team_type", "general") - org_teams_by_type[team_type] = org_teams_by_type.get(team_type, 0) + 1 - stats["team_types"][team_type] = ( - stats["team_types"].get(team_type, 0) + 1 - ) - - agents = await DatabaseManager.get_agents(team_id) - org_agent_count += len(agents) - - for agent in agents: - agent_type = agent.get("type", "developer") - org_agents_by_type[agent_type] = ( - org_agents_by_type.get(agent_type, 0) + 1 - ) - stats["agent_types"][agent_type] = ( - stats["agent_types"].get(agent_type, 0) + 1 - ) - - stats["by_organization"].append( - { - "id": org_id, - "name": org["name"], - "teams": len(teams), - "agents": org_agent_count, - "teams_by_type": org_teams_by_type, - "agents_by_type": org_agents_by_type, - } - ) - - stats["teams"] += len(teams) - stats["agents"] += org_agent_count - - return stats - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get hierarchy stats: {str(e)}" - ) - - -@router.get("/organization/{organization_id}/chart") -async def get_organization_chart(organization_id: str): - """Get detailed chart data for a specific organization""" - try: - # Get organization details - organization = await DatabaseManager.get_organization(organization_id) - if not organization: - raise HTTPException(status_code=404, detail="Organization not found") - - # Get teams and agents - teams = await DatabaseManager.get_teams(organization_id) - - chart_data = {"organization": organization, "teams": [], "total_agents": 0} - - for team in teams: - team_id = team["id"] - agents = await DatabaseManager.get_agents(team_id) - - team_data = { - "id": team_id, - "name": team["name"], - "description": team.get("description", ""), - "team_type": team.get("team_type", "general"), - "settings": team.get("settings", {}), - "agents": agents, - "agent_count": len(agents), - } - - chart_data["teams"].append(team_data) - chart_data["total_agents"] += len(agents) - - return chart_data - - except HTTPException: - raise - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get organization chart: {str(e)}" - ) - - -@router.get("/search") -async def search_hierarchy(q: str, entity_type: Optional[str] = None): - """Search across organizations, teams, and agents""" - try: - if not q or len(q) < 2: - raise HTTPException( - status_code=400, detail="Query must be at least 2 characters" - ) - - results = {"organizations": [], "teams": [], "agents": [], "total_results": 0} - - query = q.lower() - - # Search organizations - if not entity_type or entity_type == "organization": - organizations = await DatabaseManager.get_organizations() - for org in organizations: - if ( - query in org["name"].lower() - or query in org.get("description", "").lower() - ): - results["organizations"].append(org) - - # Search teams - if not entity_type or entity_type == "team": - organizations = await DatabaseManager.get_organizations() - for org in organizations: - teams = await DatabaseManager.get_teams(org["id"]) - for team in teams: - if ( - query in team["name"].lower() - or query in team.get("description", "").lower() - or query in team.get("team_type", "").lower() - ): - team["organization_name"] = org["name"] - results["teams"].append(team) - - # Search agents - if not entity_type or entity_type == "agent": - organizations = await DatabaseManager.get_organizations() - for org in organizations: - teams = await DatabaseManager.get_teams(org["id"]) - for team in teams: - agents = await DatabaseManager.get_agents(team["id"]) - for agent in agents: - if ( - query in agent["name"].lower() - or query in agent.get("role", "").lower() - or query in agent.get("type", "").lower() - ): - agent["organization_name"] = org["name"] - agent["team_name"] = team["name"] - results["agents"].append(agent) - - results["total_results"] = ( - len(results["organizations"]) - + len(results["teams"]) - + len(results["agents"]) - ) - - return results - - except HTTPException: - raise - except Exception as e: - raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}") +import asyncio +import json +from typing import Any, Dict, List, Optional + +import httpx +from fastapi import APIRouter, HTTPException + +from database import DatabaseManager + +router = APIRouter(prefix="/hierarchy", tags=["hierarchy"]) + + +@router.get("/visualization") +async def get_hierarchy_visualization(): + """ + Get complete organizational hierarchy for visualization with ReactFlow/GoJS + + Returns structured data optimized for hierarchical visualization libraries: + - Organizations as root nodes + - Teams as intermediate nodes + - Agents as leaf nodes + - Relationships and positioning data + """ + try: + # Get data from hierarchy API + async with httpx.AsyncClient() as client: + # Get all organizations + orgs_response = await client.get("http://localhost:8006/organizations") + organizations = orgs_response.json() + + # Get all teams + teams_response = await client.get("http://localhost:8006/teams") + all_teams = teams_response.json() + + # Get all agents + agents_response = await client.get("http://localhost:8006/agents") + all_agents = agents_response.json() + + if not organizations: + return {"nodes": [], "edges": [], "message": "No organizations found"} + + nodes = [] + edges = [] + y_offset = 0 + + for org_idx, org in enumerate(organizations): + org_id = org["id"] + org_node_id = f"org-{org_id}" + + # Add organization node + nodes.append( + { + "id": org_node_id, + "type": "organization", + "data": { + "label": org["name"], + "description": org.get("description", ""), + "type": "Organization", + "settings": org.get("settings", {}), + "entity_id": org_id, + }, + "position": {"x": org_idx * 800, "y": y_offset}, + "style": { + "background": "#1e40af", + "color": "white", + "border": "2px solid #1e3a8a", + "borderRadius": "12px", + "padding": "12px", + "minWidth": "200px", + }, + } + ) + + # Filter teams for this organization + teams = [t for t in all_teams if t.get("organization_id") == org_id] + team_y_offset = y_offset + 150 + + for team_idx, team in enumerate(teams): + team_id = team["id"] + team_node_id = f"team-{team_id}" + + # Add team node + nodes.append( + { + "id": team_node_id, + "type": "team", + "data": { + "label": team["name"], + "description": team.get("description", ""), + "type": f"Team ({team.get('team_type', 'general')})", + "settings": team.get("settings", {}), + "entity_id": team_id, + "organization_id": org_id, + }, + "position": { + "x": org_idx * 800 + (team_idx % 3) * 250 - 250, + "y": team_y_offset + (team_idx // 3) * 120, + }, + "style": { + "background": "#059669", + "color": "white", + "border": "2px solid #047857", + "borderRadius": "8px", + "padding": "10px", + "minWidth": "180px", + }, + } + ) + + # Add edge from organization to team + edges.append( + { + "id": f"edge-{org_node_id}-{team_node_id}", + "source": org_node_id, + "target": team_node_id, + "type": "smoothstep", + "style": {"stroke": "#64748b", "strokeWidth": 2}, + "markerEnd": {"type": "arrowclosed", "color": "#64748b"}, + } + ) + + # Filter agents for this team - note: need to check team_id in agents + agents = [ + a + for a in all_agents + if a.get("team_id") == team_id + or (hasattr(a, "config") and a.config.get("team_id") == team_id) + ] + # Fallback: if no team_id in agent data, we'll have limited agents + if not agents and all_agents: + # For now, include some agents if they don't have team assignments + agents = all_agents[:2] # Include first 2 agents as examples + + agent_y_offset = team_y_offset + (team_idx // 3) * 120 + 100 + + for agent_idx, agent in enumerate(agents): + agent_id = agent["id"] + agent_node_id = f"agent-{agent_id}" + + # Determine agent color by type + agent_colors = { + "executive": {"bg": "#dc2626", "border": "#b91c1c"}, + "developer": {"bg": "#2563eb", "border": "#1d4ed8"}, + "marketing": {"bg": "#7c3aed", "border": "#6d28d9"}, + "sales": {"bg": "#ea580c", "border": "#c2410c"}, + "qa": {"bg": "#16a34a", "border": "#15803d"}, + "devops": {"bg": "#0891b2", "border": "#0e7490"}, + "designer": {"bg": "#e11d48", "border": "#be185d"}, + } + + agent_type = agent.get("type", "developer") + colors = agent_colors.get( + agent_type, {"bg": "#6b7280", "border": "#4b5563"} + ) + + # Add agent node + nodes.append( + { + "id": agent_node_id, + "type": "agent", + "data": { + "label": agent["name"], + "role": agent.get("role", ""), + "type": f"Agent ({agent_type})", + "status": agent.get("status", "active"), + "config": agent.get("config", {}), + "entity_id": agent_id, + "team_id": team_id, + "organization_id": org_id, + }, + "position": { + "x": org_idx * 800 + + (team_idx % 3) * 250 + - 250 + + (agent_idx % 2) * 120 + - 60, + "y": agent_y_offset + (agent_idx // 2) * 80, + }, + "style": { + "background": colors["bg"], + "color": "white", + "border": f"2px solid {colors['border']}", + "borderRadius": "6px", + "padding": "8px", + "minWidth": "150px", + }, + } + ) + + # Add edge from team to agent + edges.append( + { + "id": f"edge-{team_node_id}-{agent_node_id}", + "source": team_node_id, + "target": agent_node_id, + "type": "smoothstep", + "style": {"stroke": "#94a3b8", "strokeWidth": 1.5}, + "markerEnd": {"type": "arrowclosed", "color": "#94a3b8"}, + } + ) + + return { + "nodes": nodes, + "edges": edges, + "metadata": { + "total_organizations": len(organizations), + "total_teams": len(all_teams), + "total_agents": len(all_agents), + "generated_at": "2025-08-06T10:00:00Z", + }, + } + + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"Failed to generate hierarchy visualization: {str(e)}", + ) + + +@router.get("/stats") +async def get_hierarchy_stats(): + """Get comprehensive hierarchy statistics""" + try: + organizations = await DatabaseManager.get_organizations() + + if not organizations: + return {"organizations": 0, "teams": 0, "agents": 0, "by_organization": []} + + stats = { + "organizations": len(organizations), + "teams": 0, + "agents": 0, + "by_organization": [], + "agent_types": {}, + "team_types": {}, + } + + for org in organizations: + org_id = org["id"] + teams = await DatabaseManager.get_teams(org_id) + org_agent_count = 0 + org_teams_by_type = {} + org_agents_by_type = {} + + for team in teams: + team_id = team["id"] + team_type = team.get("team_type", "general") + org_teams_by_type[team_type] = org_teams_by_type.get(team_type, 0) + 1 + stats["team_types"][team_type] = ( + stats["team_types"].get(team_type, 0) + 1 + ) + + agents = await DatabaseManager.get_agents(team_id) + org_agent_count += len(agents) + + for agent in agents: + agent_type = agent.get("type", "developer") + org_agents_by_type[agent_type] = ( + org_agents_by_type.get(agent_type, 0) + 1 + ) + stats["agent_types"][agent_type] = ( + stats["agent_types"].get(agent_type, 0) + 1 + ) + + stats["by_organization"].append( + { + "id": org_id, + "name": org["name"], + "teams": len(teams), + "agents": org_agent_count, + "teams_by_type": org_teams_by_type, + "agents_by_type": org_agents_by_type, + } + ) + + stats["teams"] += len(teams) + stats["agents"] += org_agent_count + + return stats + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get hierarchy stats: {str(e)}" + ) + + +@router.get("/organization/{organization_id}/chart") +async def get_organization_chart(organization_id: str): + """Get detailed chart data for a specific organization""" + try: + # Get organization details + organization = await DatabaseManager.get_organization(organization_id) + if not organization: + raise HTTPException(status_code=404, detail="Organization not found") + + # Get teams and agents + teams = await DatabaseManager.get_teams(organization_id) + + chart_data = {"organization": organization, "teams": [], "total_agents": 0} + + for team in teams: + team_id = team["id"] + agents = await DatabaseManager.get_agents(team_id) + + team_data = { + "id": team_id, + "name": team["name"], + "description": team.get("description", ""), + "team_type": team.get("team_type", "general"), + "settings": team.get("settings", {}), + "agents": agents, + "agent_count": len(agents), + } + + chart_data["teams"].append(team_data) + chart_data["total_agents"] += len(agents) + + return chart_data + + except HTTPException: + raise + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get organization chart: {str(e)}" + ) + + +@router.get("/search") +async def search_hierarchy(q: str, entity_type: Optional[str] = None): + """Search across organizations, teams, and agents""" + try: + if not q or len(q) < 2: + raise HTTPException( + status_code=400, detail="Query must be at least 2 characters" + ) + + results = {"organizations": [], "teams": [], "agents": [], "total_results": 0} + + query = q.lower() + + # Search organizations + if not entity_type or entity_type == "organization": + organizations = await DatabaseManager.get_organizations() + for org in organizations: + if ( + query in org["name"].lower() + or query in org.get("description", "").lower() + ): + results["organizations"].append(org) + + # Search teams + if not entity_type or entity_type == "team": + organizations = await DatabaseManager.get_organizations() + for org in organizations: + teams = await DatabaseManager.get_teams(org["id"]) + for team in teams: + if ( + query in team["name"].lower() + or query in team.get("description", "").lower() + or query in team.get("team_type", "").lower() + ): + team["organization_name"] = org["name"] + results["teams"].append(team) + + # Search agents + if not entity_type or entity_type == "agent": + organizations = await DatabaseManager.get_organizations() + for org in organizations: + teams = await DatabaseManager.get_teams(org["id"]) + for team in teams: + agents = await DatabaseManager.get_agents(team["id"]) + for agent in agents: + if ( + query in agent["name"].lower() + or query in agent.get("role", "").lower() + or query in agent.get("type", "").lower() + ): + agent["organization_name"] = org["name"] + agent["team_name"] = team["name"] + results["agents"].append(agent) + + results["total_results"] = ( + len(results["organizations"]) + + len(results["teams"]) + + len(results["agents"]) + ) + + return results + + except HTTPException: + raise + except Exception as e: + raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}") diff --git a/services/orchestrator/knowledge_propagation_engine.py b/services/orchestrator/knowledge_propagation_engine.py index 056b6f5..094b676 100644 --- a/services/orchestrator/knowledge_propagation_engine.py +++ b/services/orchestrator/knowledge_propagation_engine.py @@ -1,958 +1,958 @@ -""" -Knowledge Propagation Engine for FuzeAgent - -This module handles automated knowledge flow between agents, teams, and organizations. -It determines when knowledge should be propagated, executes the propagation, -and manages the lifecycle of knowledge across hierarchical levels. -""" - -import asyncio -import json -import logging -import uuid -from dataclasses import dataclass -from datetime import datetime, timedelta -from enum import Enum -from typing import Any, Dict, List, Optional, Set, Tuple - -import asyncpg -from sentence_transformers import SentenceTransformer - -from .organization_rag_manager import ( - ContentType, - KnowledgeCategory, - OrganizationRAGManager, - SourceType, - VisibilityLevel, -) -from .team_knowledge_manager import TeamKnowledgeManager - -logger = logging.getLogger(__name__) - - -class PropagationTrigger(str, Enum): - TASK_COMPLETION = "task_completion" - KNOWLEDGE_THRESHOLD = "knowledge_threshold" - MANUAL_REQUEST = "manual_request" - SCHEDULED_SYNC = "scheduled_sync" - CROSS_TEAM_REQUEST = "cross_team_request" - QUALITY_IMPROVEMENT = "quality_improvement" - - -class PropagationStatus(str, Enum): - PENDING = "pending" - PROCESSING = "processing" - COMPLETED = "completed" - FAILED = "failed" - REJECTED = "rejected" - - -class AcceptanceStatus(str, Enum): - PENDING = "pending" - ACCEPTED = "accepted" - REJECTED = "rejected" - MODIFIED = "modified" - - -@dataclass -class PropagationRule: - """Defines rules for knowledge propagation""" - - source_type: str # 'agent', 'team', 'organization' - target_type: str # 'agent', 'team', 'organization' - min_confidence: float - min_success_correlation: float - min_usage_count: int - knowledge_categories: List[KnowledgeCategory] - auto_approve: bool - propagation_weight: float - - -@dataclass -class PropagationTask: - """Represents a knowledge propagation task""" - - id: str - source_type: str - source_id: str - target_type: str - target_id: str - knowledge_type: str - knowledge_content_id: str - propagation_method: str - propagation_trigger: PropagationTrigger - confidence_score: float - propagation_status: PropagationStatus - acceptance_status: AcceptanceStatus - metadata: Dict[str, Any] - created_at: datetime - processed_at: Optional[datetime] - completed_at: Optional[datetime] - - -class KnowledgePropagationEngine: - """ - Manages the automated flow of knowledge across the organization hierarchy. - Handles agent → team → organization propagation and cross-team sharing. - """ - - def __init__( - self, - database_url: str, - org_rag_manager: OrganizationRAGManager, - team_knowledge_manager: TeamKnowledgeManager, - ): - self.database_url = database_url - self.org_rag_manager = org_rag_manager - self.team_knowledge_manager = team_knowledge_manager - self.pool: Optional[asyncpg.Pool] = None - - # Initialize embedding model for similarity analysis - self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") - - # Default propagation rules - self.default_rules = self._create_default_propagation_rules() - - # Configuration - self.propagation_batch_size = 50 - self.max_concurrent_propagations = 5 - self.similarity_threshold = 0.8 - self.propagation_cooldown_hours = 24 - - # Statistics - self.propagations_processed = 0 - self.propagations_completed = 0 - self.propagations_rejected = 0 - - # Background task management - self._propagation_task: Optional[asyncio.Task] = None - self._running = False - - async def initialize(self): - """Initialize the knowledge propagation engine""" - logger.info("Initializing KnowledgePropagationEngine") - - try: - self.pool = await asyncpg.create_pool( - self.database_url, min_size=2, max_size=10, command_timeout=60 - ) - - # Start background propagation processing - self._running = True - self._propagation_task = asyncio.create_task( - self._background_propagation_processor() - ) - - logger.info("KnowledgePropagationEngine initialized successfully") - - except Exception as e: - logger.error(f"Failed to initialize KnowledgePropagationEngine: {e}") - raise - - async def close(self): - """Close the propagation engine and cleanup resources""" - self._running = False - - if self._propagation_task: - self._propagation_task.cancel() - try: - await self._propagation_task - except asyncio.CancelledError: - pass - - if self.pool: - await self.pool.close() - - logger.info("KnowledgePropagationEngine closed") - - async def trigger_agent_to_team_propagation( - self, agent_id: str, task_id: str, task_outcome: Dict[str, Any] - ) -> List[str]: - """Trigger knowledge propagation from agent to team level after task completion""" - - propagation_ids = [] - - async with self.pool.acquire() as conn: - # Get agent's team - team_id = await conn.fetchval( - """ - SELECT team_id FROM agents WHERE id = $1 - """, - agent_id, - ) - - if not team_id: - logger.warning(f"No team found for agent {agent_id}") - return propagation_ids - - # Get recent agent memories from this task - recent_memories = await conn.fetch( - """ - SELECT * FROM agent_memory - WHERE agent_id = $1 - AND task_id = $2 - AND confidence_score >= 0.6 - AND propagated_to_team = FALSE - ORDER BY confidence_score DESC, created_at DESC - """, - agent_id, - task_id, - ) - - # Group memories by type and analyze for propagation - memory_groups = self._group_memories_for_propagation(recent_memories) - - for group_type, memories in memory_groups.items(): - if len(memories) >= 1 and self._meets_propagation_criteria( - memories, task_outcome - ): - # Create propagation task - propagation_id = await self._create_propagation_task( - source_type="agent", - source_id=agent_id, - target_type="team", - target_id=str(team_id), - knowledge_type=group_type, - knowledge_content_ids=[str(mem["id"]) for mem in memories], - propagation_trigger=PropagationTrigger.TASK_COMPLETION, - confidence_score=self._calculate_group_confidence(memories), - metadata={ - "task_id": task_id, - "task_outcome": task_outcome, - "memory_count": len(memories), - }, - ) - - propagation_ids.append(propagation_id) - - logger.info( - f"Created {len(propagation_ids)} propagation tasks for agent {agent_id} → team {team_id}" - ) - return propagation_ids - - async def trigger_team_to_org_propagation( - self, team_id: str, knowledge_threshold_check: bool = True - ) -> List[str]: - """Trigger knowledge propagation from team to organization level""" - - propagation_ids = [] - - async with self.pool.acquire() as conn: - # Get organization ID - org_id = await conn.fetchval( - """ - SELECT organization_id FROM teams WHERE id = $1 - """, - team_id, - ) - - if not org_id: - logger.warning(f"No organization found for team {team_id}") - return propagation_ids - - # Find high-value team knowledge for propagation - if knowledge_threshold_check: - team_knowledge = await conn.fetch( - """ - SELECT * FROM team_knowledge_base - WHERE team_id = $1 - AND effectiveness_score >= 0.7 - AND agent_adoption_rate >= 0.5 - AND created_at <= NOW() - INTERVAL '7 days' -- Allow time for validation - ORDER BY effectiveness_score DESC, agent_adoption_rate DESC - """, - team_id, - ) - else: - team_knowledge = await conn.fetch( - """ - SELECT * FROM team_knowledge_base - WHERE team_id = $1 - ORDER BY effectiveness_score DESC - LIMIT 10 - """, - team_id, - ) - - for knowledge in team_knowledge: - # Check if similar knowledge already exists at org level - if not await self._check_for_similar_org_knowledge( - knowledge, str(org_id) - ): - # Create propagation task - propagation_id = await self._create_propagation_task( - source_type="team", - source_id=team_id, - target_type="organization", - target_id=str(org_id), - knowledge_type=knowledge["knowledge_category"], - knowledge_content_ids=[str(knowledge["id"])], - propagation_trigger=PropagationTrigger.KNOWLEDGE_THRESHOLD, - confidence_score=knowledge["effectiveness_score"], - metadata={ - "team_knowledge_id": str(knowledge["id"]), - "adoption_rate": knowledge["agent_adoption_rate"], - "contributing_agents": knowledge["contributing_agents"], - }, - ) - - propagation_ids.append(propagation_id) - - logger.info( - f"Created {len(propagation_ids)} propagation tasks for team {team_id} → organization {org_id}" - ) - return propagation_ids - - async def trigger_cross_team_sharing( - self, - source_team_id: str, - knowledge_categories: List[KnowledgeCategory], - target_teams: Optional[List[str]] = None, - ) -> List[str]: - """Trigger knowledge sharing between teams""" - - propagation_ids = [] - - async with self.pool.acquire() as conn: - # Get organization and determine target teams - org_id = await conn.fetchval( - """ - SELECT organization_id FROM teams WHERE id = $1 - """, - source_team_id, - ) - - if not org_id: - return propagation_ids - - if not target_teams: - # Get all teams in the organization except source team - target_teams_rows = await conn.fetch( - """ - SELECT id FROM teams - WHERE organization_id = $1 AND id != $2 - """, - org_id, - source_team_id, - ) - target_teams = [str(row["id"]) for row in target_teams_rows] - - # Get relevant knowledge from source team - category_list = [cat.value for cat in knowledge_categories] - source_knowledge = await conn.fetch( - """ - SELECT * FROM team_knowledge_base - WHERE team_id = $1 - AND knowledge_category = ANY($2) - AND effectiveness_score >= 0.6 - ORDER BY effectiveness_score DESC - LIMIT 20 - """, - source_team_id, - category_list, - ) - - # Create propagation tasks for each target team - for target_team_id in target_teams: - for knowledge in source_knowledge: - # Check if target team would benefit from this knowledge - relevance = await self._calculate_cross_team_relevance( - knowledge, source_team_id, target_team_id - ) - - if relevance >= 0.5: - propagation_id = await self._create_propagation_task( - source_type="team", - source_id=source_team_id, - target_type="team", - target_id=target_team_id, - knowledge_type=knowledge["knowledge_category"], - knowledge_content_ids=[str(knowledge["id"])], - propagation_trigger=PropagationTrigger.CROSS_TEAM_REQUEST, - confidence_score=relevance, - metadata={ - "cross_team_relevance": relevance, - "source_effectiveness": knowledge[ - "effectiveness_score" - ], - }, - ) - - propagation_ids.append(propagation_id) - - logger.info(f"Created {len(propagation_ids)} cross-team propagation tasks") - return propagation_ids - - async def process_pending_propagations(self, limit: int = 10) -> Dict[str, int]: - """Process pending propagation tasks""" - - results = {"processed": 0, "completed": 0, "failed": 0} - - async with self.pool.acquire() as conn: - # Get pending propagation tasks - pending_tasks = await conn.fetch( - """ - SELECT * FROM knowledge_propagation_log - WHERE propagation_status = 'pending' - ORDER BY propagated_at ASC - LIMIT $1 - """, - limit, - ) - - for task_row in pending_tasks: - task = self._row_to_propagation_task(task_row) - - try: - # Update status to processing - await conn.execute( - """ - UPDATE knowledge_propagation_log - SET propagation_status = 'processing', processed_at = NOW() - WHERE id = $1 - """, - task.id, - ) - - # Process the propagation - success = await self._execute_propagation(task) - - if success: - # Mark as completed - await conn.execute( - """ - UPDATE knowledge_propagation_log - SET propagation_status = 'completed', - acceptance_status = 'accepted', - completed_at = NOW() - WHERE id = $1 - """, - task.id, - ) - results["completed"] += 1 - self.propagations_completed += 1 - else: - # Mark as failed - await conn.execute( - """ - UPDATE knowledge_propagation_log - SET propagation_status = 'failed' - WHERE id = $1 - """, - task.id, - ) - results["failed"] += 1 - - results["processed"] += 1 - self.propagations_processed += 1 - - except Exception as e: - logger.error(f"Error processing propagation task {task.id}: {e}") - await conn.execute( - """ - UPDATE knowledge_propagation_log - SET propagation_status = 'failed', - metadata = metadata || $2 - WHERE id = $1 - """, - task.id, - json.dumps({"error": str(e)}), - ) - results["failed"] += 1 - - return results - - async def get_propagation_statistics( - self, - organization_id: Optional[str] = None, - team_id: Optional[str] = None, - days_back: int = 30, - ) -> Dict[str, Any]: - """Get comprehensive propagation statistics""" - - async with self.pool.acquire() as conn: - where_conditions = [] - params = [] - - if organization_id: - where_conditions.append( - "target_id = $1 AND target_type = 'organization'" - ) - params.append(organization_id) - elif team_id: - where_conditions.append( - "(target_id = $1 OR source_id = $1) AND ('team' = ANY(ARRAY[target_type, source_type]))" - ) - params.append(team_id) - - # Parameterize the time window; days_back is bound, not interpolated. - days_param_idx = len(params) + 1 - params.append(days_back) - where_conditions.append( - f"propagated_at >= NOW() - (INTERVAL '1 day' * ${days_param_idx})" - ) - - where_clause = "WHERE " + " AND ".join(where_conditions) - - # Basic statistics - stats = await conn.fetchrow( - f""" - SELECT - COUNT(*) as total_propagations, - COUNT(CASE WHEN propagation_status = 'completed' THEN 1 END) as completed, - COUNT(CASE WHEN propagation_status = 'failed' THEN 1 END) as failed, - COUNT(CASE WHEN propagation_status = 'pending' THEN 1 END) as pending, - COUNT(CASE WHEN acceptance_status = 'accepted' THEN 1 END) as accepted, - COUNT(CASE WHEN acceptance_status = 'rejected' THEN 1 END) as rejected, - AVG(confidence_score) as avg_confidence - FROM knowledge_propagation_log - {where_clause} - """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params - *params, - ) - - # Propagation flow statistics - flow_stats = await conn.fetch( - f""" - SELECT - source_type || ' → ' || target_type as flow_type, - COUNT(*) as count, - AVG(confidence_score) as avg_confidence, - COUNT(CASE WHEN propagation_status = 'completed' THEN 1 END)::float / COUNT(*) as success_rate - FROM knowledge_propagation_log - {where_clause} - GROUP BY source_type, target_type - ORDER BY count DESC - """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params - *params, - ) - - # Trigger analysis - trigger_stats = await conn.fetch( - f""" - SELECT - propagation_trigger, - COUNT(*) as count, - AVG(confidence_score) as avg_confidence - FROM knowledge_propagation_log - {where_clause} - GROUP BY propagation_trigger - ORDER BY count DESC - """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params - *params, - ) - - return { - "time_period_days": days_back, - "basic_stats": dict(stats) if stats else {}, - "flow_patterns": [dict(flow) for flow in flow_stats], - "trigger_analysis": [dict(trigger) for trigger in trigger_stats], - "generated_at": datetime.now().isoformat(), - } - - async def _background_propagation_processor(self): - """Background task to continuously process propagation queue""" - - while self._running: - try: - # Process a batch of propagations - results = await self.process_pending_propagations( - self.propagation_batch_size - ) - - if results["processed"] > 0: - logger.info( - f"Processed {results['processed']} propagations: " - f"{results['completed']} completed, {results['failed']} failed" - ) - - # Sleep between processing cycles - await asyncio.sleep(30) # Process every 30 seconds - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Error in background propagation processor: {e}") - await asyncio.sleep(60) # Wait longer after errors - - async def _create_propagation_task( - self, - source_type: str, - source_id: str, - target_type: str, - target_id: str, - knowledge_type: str, - knowledge_content_ids: List[str], - propagation_trigger: PropagationTrigger, - confidence_score: float, - metadata: Dict[str, Any], - ) -> str: - """Create a new propagation task""" - - propagation_id = str(uuid.uuid4()) - - async with self.pool.acquire() as conn: - await conn.execute( - """ - INSERT INTO knowledge_propagation_log ( - id, source_type, source_id, target_type, target_id, - knowledge_type, propagation_method, propagation_trigger, - confidence_score, propagation_status, acceptance_status, metadata - ) VALUES ($1, $2, $3, $4, $5, $6, 'automatic', $7, $8, 'pending', 'pending', $9) - """, - propagation_id, - source_type, - source_id, - target_type, - target_id, - knowledge_type, - propagation_trigger.value, - confidence_score, - json.dumps({**metadata, "content_ids": knowledge_content_ids}), - ) - - return propagation_id - - async def _execute_propagation(self, task: PropagationTask) -> bool: - """Execute a specific propagation task""" - - try: - if task.source_type == "agent" and task.target_type == "team": - return await self._execute_agent_to_team_propagation(task) - elif task.source_type == "team" and task.target_type == "organization": - return await self._execute_team_to_org_propagation(task) - elif task.source_type == "team" and task.target_type == "team": - return await self._execute_team_to_team_propagation(task) - else: - logger.warning( - f"Unsupported propagation type: {task.source_type} → {task.target_type}" - ) - return False - - except Exception as e: - logger.error(f"Error executing propagation {task.id}: {e}") - return False - - async def _execute_agent_to_team_propagation(self, task: PropagationTask) -> bool: - """Execute agent → team propagation""" - - content_ids = task.metadata.get("content_ids", []) - if not content_ids: - return False - - # Aggregate agent memories to team knowledge - result = await self.team_knowledge_manager.aggregate_agent_knowledge_to_team( - team_id=task.target_id, - agent_id=task.source_id, - agent_memory_ids=content_ids, - aggregation_method="propagation", - ) - - return result is not None - - async def _execute_team_to_org_propagation(self, task: PropagationTask) -> bool: - """Execute team → organization propagation""" - - team_knowledge_id = task.metadata.get("team_knowledge_id") - if not team_knowledge_id: - return False - - async with self.pool.acquire() as conn: - # Get team knowledge - team_knowledge = await conn.fetchrow( - """ - SELECT * FROM team_knowledge_base WHERE id = $1 - """, - team_knowledge_id, - ) - - if not team_knowledge: - return False - - # Create organization knowledge - org_knowledge_id = await self.org_rag_manager.add_knowledge( - organization_id=task.target_id, - title=f"[Team Contribution] {team_knowledge['title']}", - content=team_knowledge["content"], - content_type=ContentType(team_knowledge["content_type"]), - knowledge_category=KnowledgeCategory( - team_knowledge["knowledge_category"] - ), - source_type=SourceType.TEAM_AGGREGATION, - source_team_id=task.source_id, - relevance_score=team_knowledge["effectiveness_score"], - quality_score=team_knowledge["effectiveness_score"], - visibility_level=VisibilityLevel.ORGANIZATION, - metadata={ - "team_adoption_rate": team_knowledge["agent_adoption_rate"], - "team_effectiveness": team_knowledge["effectiveness_score"], - "contributing_agents": team_knowledge["contributing_agents"], - "propagation_task_id": task.id, - }, - tags=team_knowledge["tags"] + ["team_contribution"], - ) - - return org_knowledge_id is not None - - async def _execute_team_to_team_propagation(self, task: PropagationTask) -> bool: - """Execute team → team propagation""" - - team_knowledge_id = task.metadata.get("content_ids", [None])[0] - if not team_knowledge_id: - return False - - async with self.pool.acquire() as conn: - # Get source team knowledge - source_knowledge = await conn.fetchrow( - """ - SELECT * FROM team_knowledge_base WHERE id = $1 - """, - team_knowledge_id, - ) - - if not source_knowledge: - return False - - # Create adapted knowledge for target team - adapted_knowledge_id = ( - await self.team_knowledge_manager.create_team_knowledge( - team_id=task.target_id, - title=f"[Shared] {source_knowledge['title']}", - content=source_knowledge["content"], - content_type=ContentType(source_knowledge["content_type"]), - knowledge_category=KnowledgeCategory( - source_knowledge["knowledge_category"] - ), - source_type=SourceType.TEAM_AGGREGATION, - contributing_agents=[], - source_knowledge_ids=[team_knowledge_id], - aggregation_method="cross_team_sharing", - team_relevance_score=task.confidence_score, - metadata={ - "source_team_id": task.source_id, - "cross_team_propagation": True, - "original_effectiveness": source_knowledge[ - "effectiveness_score" - ], - "propagation_task_id": task.id, - }, - tags=source_knowledge["tags"] + ["cross_team_shared"], - ) - ) - - return adapted_knowledge_id is not None - - def _group_memories_for_propagation(self, memories: List) -> Dict[str, List]: - """Group memories by type for propagation analysis""" - groups = {} - - for memory in memories: - memory_type = memory.get("memory_type", "general") - if memory_type not in groups: - groups[memory_type] = [] - groups[memory_type].append(memory) - - return groups - - def _meets_propagation_criteria( - self, memories: List, task_outcome: Dict[str, Any] - ) -> bool: - """Determine if memories meet criteria for propagation""" - - # Check success rate - success = task_outcome.get("success", False) - if not success: - return False - - # Check confidence scores - avg_confidence = sum(mem["confidence_score"] for mem in memories) / len( - memories - ) - if avg_confidence < 0.6: - return False - - # Check memory age (don't propagate very old memories) - recent_memories = [ - mem for mem in memories if (datetime.now() - mem["created_at"]).days <= 7 - ] - - return ( - len(recent_memories) >= len(memories) * 0.5 - ) # At least 50% should be recent - - def _calculate_group_confidence(self, memories: List) -> float: - """Calculate confidence score for a group of memories""" - if not memories: - return 0.0 - - scores = [mem["confidence_score"] for mem in memories] - success_correlations = [mem.get("success_correlation", 0.0) for mem in memories] - - # Weighted average with recency bias - weights = [1.0 / (1 + i * 0.1) for i in range(len(memories))] - - weighted_confidence = sum(s * w for s, w in zip(scores, weights)) / sum(weights) - avg_success = ( - sum(success_correlations) / len(success_correlations) - if success_correlations - else 0.0 - ) - - return min(1.0, weighted_confidence * 0.7 + avg_success * 0.3) - - async def _check_for_similar_org_knowledge( - self, team_knowledge: Dict, org_id: str - ) -> bool: - """Check if similar knowledge already exists at organization level""" - - if not team_knowledge.get("embedding"): - return False - - # Search for similar content - search_results = await self.org_rag_manager.search_knowledge( - organization_id=org_id, - query=team_knowledge["content"][:200], # Use beginning of content as query - limit=5, - min_similarity=self.similarity_threshold, - ) - - # Check if any results are highly similar - for result in search_results: - if result.similarity_score >= self.similarity_threshold: - return True - - return False - - async def _calculate_cross_team_relevance( - self, knowledge: Dict, source_team_id: str, target_team_id: str - ) -> float: - """Calculate how relevant knowledge from one team is for another team""" - - async with self.pool.acquire() as conn: - # Get team information - teams = await conn.fetch( - """ - SELECT id, team_type, settings FROM teams - WHERE id IN ($1, $2) - """, - source_team_id, - target_team_id, - ) - - if len(teams) != 2: - return 0.0 - - source_team = next(t for t in teams if str(t["id"]) == source_team_id) - target_team = next(t for t in teams if str(t["id"]) == target_team_id) - - relevance_factors = [] - - # Factor 1: Team type similarity - type_similarity = ( - 1.0 if source_team["team_type"] == target_team["team_type"] else 0.3 - ) - relevance_factors.append(type_similarity) - - # Factor 2: Knowledge category relevance to target team - target_team_categories = await conn.fetch( - """ - SELECT knowledge_category, COUNT(*) as usage_count - FROM team_knowledge_base - WHERE team_id = $1 - GROUP BY knowledge_category - ORDER BY usage_count DESC - LIMIT 5 - """, - target_team_id, - ) - - target_categories = [ - cat["knowledge_category"] for cat in target_team_categories - ] - category_relevance = ( - 1.0 if knowledge["knowledge_category"] in target_categories else 0.4 - ) - relevance_factors.append(category_relevance) - - # Factor 3: Effectiveness score of source knowledge - effectiveness_factor = knowledge["effectiveness_score"] - relevance_factors.append(effectiveness_factor) - - # Factor 4: Adoption rate in source team (indicates broad utility) - adoption_factor = knowledge["agent_adoption_rate"] - relevance_factors.append(adoption_factor) - - # Calculate weighted relevance - weights = [0.2, 0.3, 0.3, 0.2] - relevance = sum(f * w for f, w in zip(relevance_factors, weights)) - - return min(1.0, max(0.0, relevance)) - - def _create_default_propagation_rules(self) -> List[PropagationRule]: - """Create default propagation rules""" - return [ - # Agent to Team rules - PropagationRule( - source_type="agent", - target_type="team", - min_confidence=0.6, - min_success_correlation=0.0, - min_usage_count=1, - knowledge_categories=list(KnowledgeCategory), - auto_approve=True, - propagation_weight=1.0, - ), - # Team to Organization rules - PropagationRule( - source_type="team", - target_type="organization", - min_confidence=0.7, - min_success_correlation=0.5, - min_usage_count=3, - knowledge_categories=list(KnowledgeCategory), - auto_approve=False, - propagation_weight=0.8, - ), - # Cross-team rules - PropagationRule( - source_type="team", - target_type="team", - min_confidence=0.65, - min_success_correlation=0.4, - min_usage_count=2, - knowledge_categories=[ - KnowledgeCategory.DEVELOPMENT, - KnowledgeCategory.TESTING, - KnowledgeCategory.TROUBLESHOOTING, - ], - auto_approve=False, - propagation_weight=0.6, - ), - ] - - def _row_to_propagation_task(self, row) -> PropagationTask: - """Convert database row to PropagationTask object""" - return PropagationTask( - id=str(row["id"]), - source_type=row["source_type"], - source_id=str(row["source_id"]), - target_type=row["target_type"], - target_id=str(row["target_id"]), - knowledge_type=row["knowledge_type"], - knowledge_content_id=( - str(row["knowledge_content_id"]) if row["knowledge_content_id"] else "" - ), - propagation_method=row["propagation_method"], - propagation_trigger=PropagationTrigger(row["propagation_trigger"]), - confidence_score=row["confidence_score"], - propagation_status=PropagationStatus(row["propagation_status"]), - acceptance_status=AcceptanceStatus(row["acceptance_status"]), - metadata=( - json.loads(row["metadata"]) - if isinstance(row["metadata"], str) - else row["metadata"] - ), - created_at=row["propagated_at"], - processed_at=row["processed_at"], - completed_at=row["completed_at"], - ) +""" +Knowledge Propagation Engine for FuzeAgent + +This module handles automated knowledge flow between agents, teams, and organizations. +It determines when knowledge should be propagated, executes the propagation, +and manages the lifecycle of knowledge across hierarchical levels. +""" + +import asyncio +import json +import logging +import uuid +from dataclasses import dataclass +from datetime import datetime, timedelta +from enum import Enum +from typing import Any, Dict, List, Optional, Set, Tuple + +import asyncpg +from sentence_transformers import SentenceTransformer + +from .organization_rag_manager import ( + ContentType, + KnowledgeCategory, + OrganizationRAGManager, + SourceType, + VisibilityLevel, +) +from .team_knowledge_manager import TeamKnowledgeManager + +logger = logging.getLogger(__name__) + + +class PropagationTrigger(str, Enum): + TASK_COMPLETION = "task_completion" + KNOWLEDGE_THRESHOLD = "knowledge_threshold" + MANUAL_REQUEST = "manual_request" + SCHEDULED_SYNC = "scheduled_sync" + CROSS_TEAM_REQUEST = "cross_team_request" + QUALITY_IMPROVEMENT = "quality_improvement" + + +class PropagationStatus(str, Enum): + PENDING = "pending" + PROCESSING = "processing" + COMPLETED = "completed" + FAILED = "failed" + REJECTED = "rejected" + + +class AcceptanceStatus(str, Enum): + PENDING = "pending" + ACCEPTED = "accepted" + REJECTED = "rejected" + MODIFIED = "modified" + + +@dataclass +class PropagationRule: + """Defines rules for knowledge propagation""" + + source_type: str # 'agent', 'team', 'organization' + target_type: str # 'agent', 'team', 'organization' + min_confidence: float + min_success_correlation: float + min_usage_count: int + knowledge_categories: List[KnowledgeCategory] + auto_approve: bool + propagation_weight: float + + +@dataclass +class PropagationTask: + """Represents a knowledge propagation task""" + + id: str + source_type: str + source_id: str + target_type: str + target_id: str + knowledge_type: str + knowledge_content_id: str + propagation_method: str + propagation_trigger: PropagationTrigger + confidence_score: float + propagation_status: PropagationStatus + acceptance_status: AcceptanceStatus + metadata: Dict[str, Any] + created_at: datetime + processed_at: Optional[datetime] + completed_at: Optional[datetime] + + +class KnowledgePropagationEngine: + """ + Manages the automated flow of knowledge across the organization hierarchy. + Handles agent → team → organization propagation and cross-team sharing. + """ + + def __init__( + self, + database_url: str, + org_rag_manager: OrganizationRAGManager, + team_knowledge_manager: TeamKnowledgeManager, + ): + self.database_url = database_url + self.org_rag_manager = org_rag_manager + self.team_knowledge_manager = team_knowledge_manager + self.pool: Optional[asyncpg.Pool] = None + + # Initialize embedding model for similarity analysis + self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") + + # Default propagation rules + self.default_rules = self._create_default_propagation_rules() + + # Configuration + self.propagation_batch_size = 50 + self.max_concurrent_propagations = 5 + self.similarity_threshold = 0.8 + self.propagation_cooldown_hours = 24 + + # Statistics + self.propagations_processed = 0 + self.propagations_completed = 0 + self.propagations_rejected = 0 + + # Background task management + self._propagation_task: Optional[asyncio.Task] = None + self._running = False + + async def initialize(self): + """Initialize the knowledge propagation engine""" + logger.info("Initializing KnowledgePropagationEngine") + + try: + self.pool = await asyncpg.create_pool( + self.database_url, min_size=2, max_size=10, command_timeout=60 + ) + + # Start background propagation processing + self._running = True + self._propagation_task = asyncio.create_task( + self._background_propagation_processor() + ) + + logger.info("KnowledgePropagationEngine initialized successfully") + + except Exception as e: + logger.error(f"Failed to initialize KnowledgePropagationEngine: {e}") + raise + + async def close(self): + """Close the propagation engine and cleanup resources""" + self._running = False + + if self._propagation_task: + self._propagation_task.cancel() + try: + await self._propagation_task + except asyncio.CancelledError: + pass + + if self.pool: + await self.pool.close() + + logger.info("KnowledgePropagationEngine closed") + + async def trigger_agent_to_team_propagation( + self, agent_id: str, task_id: str, task_outcome: Dict[str, Any] + ) -> List[str]: + """Trigger knowledge propagation from agent to team level after task completion""" + + propagation_ids = [] + + async with self.pool.acquire() as conn: + # Get agent's team + team_id = await conn.fetchval( + """ + SELECT team_id FROM agents WHERE id = $1 + """, + agent_id, + ) + + if not team_id: + logger.warning(f"No team found for agent {agent_id}") + return propagation_ids + + # Get recent agent memories from this task + recent_memories = await conn.fetch( + """ + SELECT * FROM agent_memory + WHERE agent_id = $1 + AND task_id = $2 + AND confidence_score >= 0.6 + AND propagated_to_team = FALSE + ORDER BY confidence_score DESC, created_at DESC + """, + agent_id, + task_id, + ) + + # Group memories by type and analyze for propagation + memory_groups = self._group_memories_for_propagation(recent_memories) + + for group_type, memories in memory_groups.items(): + if len(memories) >= 1 and self._meets_propagation_criteria( + memories, task_outcome + ): + # Create propagation task + propagation_id = await self._create_propagation_task( + source_type="agent", + source_id=agent_id, + target_type="team", + target_id=str(team_id), + knowledge_type=group_type, + knowledge_content_ids=[str(mem["id"]) for mem in memories], + propagation_trigger=PropagationTrigger.TASK_COMPLETION, + confidence_score=self._calculate_group_confidence(memories), + metadata={ + "task_id": task_id, + "task_outcome": task_outcome, + "memory_count": len(memories), + }, + ) + + propagation_ids.append(propagation_id) + + logger.info( + f"Created {len(propagation_ids)} propagation tasks for agent {agent_id} → team {team_id}" + ) + return propagation_ids + + async def trigger_team_to_org_propagation( + self, team_id: str, knowledge_threshold_check: bool = True + ) -> List[str]: + """Trigger knowledge propagation from team to organization level""" + + propagation_ids = [] + + async with self.pool.acquire() as conn: + # Get organization ID + org_id = await conn.fetchval( + """ + SELECT organization_id FROM teams WHERE id = $1 + """, + team_id, + ) + + if not org_id: + logger.warning(f"No organization found for team {team_id}") + return propagation_ids + + # Find high-value team knowledge for propagation + if knowledge_threshold_check: + team_knowledge = await conn.fetch( + """ + SELECT * FROM team_knowledge_base + WHERE team_id = $1 + AND effectiveness_score >= 0.7 + AND agent_adoption_rate >= 0.5 + AND created_at <= NOW() - INTERVAL '7 days' -- Allow time for validation + ORDER BY effectiveness_score DESC, agent_adoption_rate DESC + """, + team_id, + ) + else: + team_knowledge = await conn.fetch( + """ + SELECT * FROM team_knowledge_base + WHERE team_id = $1 + ORDER BY effectiveness_score DESC + LIMIT 10 + """, + team_id, + ) + + for knowledge in team_knowledge: + # Check if similar knowledge already exists at org level + if not await self._check_for_similar_org_knowledge( + knowledge, str(org_id) + ): + # Create propagation task + propagation_id = await self._create_propagation_task( + source_type="team", + source_id=team_id, + target_type="organization", + target_id=str(org_id), + knowledge_type=knowledge["knowledge_category"], + knowledge_content_ids=[str(knowledge["id"])], + propagation_trigger=PropagationTrigger.KNOWLEDGE_THRESHOLD, + confidence_score=knowledge["effectiveness_score"], + metadata={ + "team_knowledge_id": str(knowledge["id"]), + "adoption_rate": knowledge["agent_adoption_rate"], + "contributing_agents": knowledge["contributing_agents"], + }, + ) + + propagation_ids.append(propagation_id) + + logger.info( + f"Created {len(propagation_ids)} propagation tasks for team {team_id} → organization {org_id}" + ) + return propagation_ids + + async def trigger_cross_team_sharing( + self, + source_team_id: str, + knowledge_categories: List[KnowledgeCategory], + target_teams: Optional[List[str]] = None, + ) -> List[str]: + """Trigger knowledge sharing between teams""" + + propagation_ids = [] + + async with self.pool.acquire() as conn: + # Get organization and determine target teams + org_id = await conn.fetchval( + """ + SELECT organization_id FROM teams WHERE id = $1 + """, + source_team_id, + ) + + if not org_id: + return propagation_ids + + if not target_teams: + # Get all teams in the organization except source team + target_teams_rows = await conn.fetch( + """ + SELECT id FROM teams + WHERE organization_id = $1 AND id != $2 + """, + org_id, + source_team_id, + ) + target_teams = [str(row["id"]) for row in target_teams_rows] + + # Get relevant knowledge from source team + category_list = [cat.value for cat in knowledge_categories] + source_knowledge = await conn.fetch( + """ + SELECT * FROM team_knowledge_base + WHERE team_id = $1 + AND knowledge_category = ANY($2) + AND effectiveness_score >= 0.6 + ORDER BY effectiveness_score DESC + LIMIT 20 + """, + source_team_id, + category_list, + ) + + # Create propagation tasks for each target team + for target_team_id in target_teams: + for knowledge in source_knowledge: + # Check if target team would benefit from this knowledge + relevance = await self._calculate_cross_team_relevance( + knowledge, source_team_id, target_team_id + ) + + if relevance >= 0.5: + propagation_id = await self._create_propagation_task( + source_type="team", + source_id=source_team_id, + target_type="team", + target_id=target_team_id, + knowledge_type=knowledge["knowledge_category"], + knowledge_content_ids=[str(knowledge["id"])], + propagation_trigger=PropagationTrigger.CROSS_TEAM_REQUEST, + confidence_score=relevance, + metadata={ + "cross_team_relevance": relevance, + "source_effectiveness": knowledge[ + "effectiveness_score" + ], + }, + ) + + propagation_ids.append(propagation_id) + + logger.info(f"Created {len(propagation_ids)} cross-team propagation tasks") + return propagation_ids + + async def process_pending_propagations(self, limit: int = 10) -> Dict[str, int]: + """Process pending propagation tasks""" + + results = {"processed": 0, "completed": 0, "failed": 0} + + async with self.pool.acquire() as conn: + # Get pending propagation tasks + pending_tasks = await conn.fetch( + """ + SELECT * FROM knowledge_propagation_log + WHERE propagation_status = 'pending' + ORDER BY propagated_at ASC + LIMIT $1 + """, + limit, + ) + + for task_row in pending_tasks: + task = self._row_to_propagation_task(task_row) + + try: + # Update status to processing + await conn.execute( + """ + UPDATE knowledge_propagation_log + SET propagation_status = 'processing', processed_at = NOW() + WHERE id = $1 + """, + task.id, + ) + + # Process the propagation + success = await self._execute_propagation(task) + + if success: + # Mark as completed + await conn.execute( + """ + UPDATE knowledge_propagation_log + SET propagation_status = 'completed', + acceptance_status = 'accepted', + completed_at = NOW() + WHERE id = $1 + """, + task.id, + ) + results["completed"] += 1 + self.propagations_completed += 1 + else: + # Mark as failed + await conn.execute( + """ + UPDATE knowledge_propagation_log + SET propagation_status = 'failed' + WHERE id = $1 + """, + task.id, + ) + results["failed"] += 1 + + results["processed"] += 1 + self.propagations_processed += 1 + + except Exception as e: + logger.error(f"Error processing propagation task {task.id}: {e}") + await conn.execute( + """ + UPDATE knowledge_propagation_log + SET propagation_status = 'failed', + metadata = metadata || $2 + WHERE id = $1 + """, + task.id, + json.dumps({"error": str(e)}), + ) + results["failed"] += 1 + + return results + + async def get_propagation_statistics( + self, + organization_id: Optional[str] = None, + team_id: Optional[str] = None, + days_back: int = 30, + ) -> Dict[str, Any]: + """Get comprehensive propagation statistics""" + + async with self.pool.acquire() as conn: + where_conditions = [] + params = [] + + if organization_id: + where_conditions.append( + "target_id = $1 AND target_type = 'organization'" + ) + params.append(organization_id) + elif team_id: + where_conditions.append( + "(target_id = $1 OR source_id = $1) AND ('team' = ANY(ARRAY[target_type, source_type]))" + ) + params.append(team_id) + + # Parameterize the time window; days_back is bound, not interpolated. + days_param_idx = len(params) + 1 + params.append(days_back) + where_conditions.append( + f"propagated_at >= NOW() - (INTERVAL '1 day' * ${days_param_idx})" + ) + + where_clause = "WHERE " + " AND ".join(where_conditions) + + # Basic statistics + stats = await conn.fetchrow( + f""" + SELECT + COUNT(*) as total_propagations, + COUNT(CASE WHEN propagation_status = 'completed' THEN 1 END) as completed, + COUNT(CASE WHEN propagation_status = 'failed' THEN 1 END) as failed, + COUNT(CASE WHEN propagation_status = 'pending' THEN 1 END) as pending, + COUNT(CASE WHEN acceptance_status = 'accepted' THEN 1 END) as accepted, + COUNT(CASE WHEN acceptance_status = 'rejected' THEN 1 END) as rejected, + AVG(confidence_score) as avg_confidence + FROM knowledge_propagation_log + {where_clause} + """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params + *params, + ) + + # Propagation flow statistics + flow_stats = await conn.fetch( + f""" + SELECT + source_type || ' → ' || target_type as flow_type, + COUNT(*) as count, + AVG(confidence_score) as avg_confidence, + COUNT(CASE WHEN propagation_status = 'completed' THEN 1 END)::float / COUNT(*) as success_rate + FROM knowledge_propagation_log + {where_clause} + GROUP BY source_type, target_type + ORDER BY count DESC + """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params + *params, + ) + + # Trigger analysis + trigger_stats = await conn.fetch( + f""" + SELECT + propagation_trigger, + COUNT(*) as count, + AVG(confidence_score) as avg_confidence + FROM knowledge_propagation_log + {where_clause} + GROUP BY propagation_trigger + ORDER BY count DESC + """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params + *params, + ) + + return { + "time_period_days": days_back, + "basic_stats": dict(stats) if stats else {}, + "flow_patterns": [dict(flow) for flow in flow_stats], + "trigger_analysis": [dict(trigger) for trigger in trigger_stats], + "generated_at": datetime.now().isoformat(), + } + + async def _background_propagation_processor(self): + """Background task to continuously process propagation queue""" + + while self._running: + try: + # Process a batch of propagations + results = await self.process_pending_propagations( + self.propagation_batch_size + ) + + if results["processed"] > 0: + logger.info( + f"Processed {results['processed']} propagations: " + f"{results['completed']} completed, {results['failed']} failed" + ) + + # Sleep between processing cycles + await asyncio.sleep(30) # Process every 30 seconds + + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"Error in background propagation processor: {e}") + await asyncio.sleep(60) # Wait longer after errors + + async def _create_propagation_task( + self, + source_type: str, + source_id: str, + target_type: str, + target_id: str, + knowledge_type: str, + knowledge_content_ids: List[str], + propagation_trigger: PropagationTrigger, + confidence_score: float, + metadata: Dict[str, Any], + ) -> str: + """Create a new propagation task""" + + propagation_id = str(uuid.uuid4()) + + async with self.pool.acquire() as conn: + await conn.execute( + """ + INSERT INTO knowledge_propagation_log ( + id, source_type, source_id, target_type, target_id, + knowledge_type, propagation_method, propagation_trigger, + confidence_score, propagation_status, acceptance_status, metadata + ) VALUES ($1, $2, $3, $4, $5, $6, 'automatic', $7, $8, 'pending', 'pending', $9) + """, + propagation_id, + source_type, + source_id, + target_type, + target_id, + knowledge_type, + propagation_trigger.value, + confidence_score, + json.dumps({**metadata, "content_ids": knowledge_content_ids}), + ) + + return propagation_id + + async def _execute_propagation(self, task: PropagationTask) -> bool: + """Execute a specific propagation task""" + + try: + if task.source_type == "agent" and task.target_type == "team": + return await self._execute_agent_to_team_propagation(task) + elif task.source_type == "team" and task.target_type == "organization": + return await self._execute_team_to_org_propagation(task) + elif task.source_type == "team" and task.target_type == "team": + return await self._execute_team_to_team_propagation(task) + else: + logger.warning( + f"Unsupported propagation type: {task.source_type} → {task.target_type}" + ) + return False + + except Exception as e: + logger.error(f"Error executing propagation {task.id}: {e}") + return False + + async def _execute_agent_to_team_propagation(self, task: PropagationTask) -> bool: + """Execute agent → team propagation""" + + content_ids = task.metadata.get("content_ids", []) + if not content_ids: + return False + + # Aggregate agent memories to team knowledge + result = await self.team_knowledge_manager.aggregate_agent_knowledge_to_team( + team_id=task.target_id, + agent_id=task.source_id, + agent_memory_ids=content_ids, + aggregation_method="propagation", + ) + + return result is not None + + async def _execute_team_to_org_propagation(self, task: PropagationTask) -> bool: + """Execute team → organization propagation""" + + team_knowledge_id = task.metadata.get("team_knowledge_id") + if not team_knowledge_id: + return False + + async with self.pool.acquire() as conn: + # Get team knowledge + team_knowledge = await conn.fetchrow( + """ + SELECT * FROM team_knowledge_base WHERE id = $1 + """, + team_knowledge_id, + ) + + if not team_knowledge: + return False + + # Create organization knowledge + org_knowledge_id = await self.org_rag_manager.add_knowledge( + organization_id=task.target_id, + title=f"[Team Contribution] {team_knowledge['title']}", + content=team_knowledge["content"], + content_type=ContentType(team_knowledge["content_type"]), + knowledge_category=KnowledgeCategory( + team_knowledge["knowledge_category"] + ), + source_type=SourceType.TEAM_AGGREGATION, + source_team_id=task.source_id, + relevance_score=team_knowledge["effectiveness_score"], + quality_score=team_knowledge["effectiveness_score"], + visibility_level=VisibilityLevel.ORGANIZATION, + metadata={ + "team_adoption_rate": team_knowledge["agent_adoption_rate"], + "team_effectiveness": team_knowledge["effectiveness_score"], + "contributing_agents": team_knowledge["contributing_agents"], + "propagation_task_id": task.id, + }, + tags=team_knowledge["tags"] + ["team_contribution"], + ) + + return org_knowledge_id is not None + + async def _execute_team_to_team_propagation(self, task: PropagationTask) -> bool: + """Execute team → team propagation""" + + team_knowledge_id = task.metadata.get("content_ids", [None])[0] + if not team_knowledge_id: + return False + + async with self.pool.acquire() as conn: + # Get source team knowledge + source_knowledge = await conn.fetchrow( + """ + SELECT * FROM team_knowledge_base WHERE id = $1 + """, + team_knowledge_id, + ) + + if not source_knowledge: + return False + + # Create adapted knowledge for target team + adapted_knowledge_id = ( + await self.team_knowledge_manager.create_team_knowledge( + team_id=task.target_id, + title=f"[Shared] {source_knowledge['title']}", + content=source_knowledge["content"], + content_type=ContentType(source_knowledge["content_type"]), + knowledge_category=KnowledgeCategory( + source_knowledge["knowledge_category"] + ), + source_type=SourceType.TEAM_AGGREGATION, + contributing_agents=[], + source_knowledge_ids=[team_knowledge_id], + aggregation_method="cross_team_sharing", + team_relevance_score=task.confidence_score, + metadata={ + "source_team_id": task.source_id, + "cross_team_propagation": True, + "original_effectiveness": source_knowledge[ + "effectiveness_score" + ], + "propagation_task_id": task.id, + }, + tags=source_knowledge["tags"] + ["cross_team_shared"], + ) + ) + + return adapted_knowledge_id is not None + + def _group_memories_for_propagation(self, memories: List) -> Dict[str, List]: + """Group memories by type for propagation analysis""" + groups = {} + + for memory in memories: + memory_type = memory.get("memory_type", "general") + if memory_type not in groups: + groups[memory_type] = [] + groups[memory_type].append(memory) + + return groups + + def _meets_propagation_criteria( + self, memories: List, task_outcome: Dict[str, Any] + ) -> bool: + """Determine if memories meet criteria for propagation""" + + # Check success rate + success = task_outcome.get("success", False) + if not success: + return False + + # Check confidence scores + avg_confidence = sum(mem["confidence_score"] for mem in memories) / len( + memories + ) + if avg_confidence < 0.6: + return False + + # Check memory age (don't propagate very old memories) + recent_memories = [ + mem for mem in memories if (datetime.now() - mem["created_at"]).days <= 7 + ] + + return ( + len(recent_memories) >= len(memories) * 0.5 + ) # At least 50% should be recent + + def _calculate_group_confidence(self, memories: List) -> float: + """Calculate confidence score for a group of memories""" + if not memories: + return 0.0 + + scores = [mem["confidence_score"] for mem in memories] + success_correlations = [mem.get("success_correlation", 0.0) for mem in memories] + + # Weighted average with recency bias + weights = [1.0 / (1 + i * 0.1) for i in range(len(memories))] + + weighted_confidence = sum(s * w for s, w in zip(scores, weights)) / sum(weights) + avg_success = ( + sum(success_correlations) / len(success_correlations) + if success_correlations + else 0.0 + ) + + return min(1.0, weighted_confidence * 0.7 + avg_success * 0.3) + + async def _check_for_similar_org_knowledge( + self, team_knowledge: Dict, org_id: str + ) -> bool: + """Check if similar knowledge already exists at organization level""" + + if not team_knowledge.get("embedding"): + return False + + # Search for similar content + search_results = await self.org_rag_manager.search_knowledge( + organization_id=org_id, + query=team_knowledge["content"][:200], # Use beginning of content as query + limit=5, + min_similarity=self.similarity_threshold, + ) + + # Check if any results are highly similar + for result in search_results: + if result.similarity_score >= self.similarity_threshold: + return True + + return False + + async def _calculate_cross_team_relevance( + self, knowledge: Dict, source_team_id: str, target_team_id: str + ) -> float: + """Calculate how relevant knowledge from one team is for another team""" + + async with self.pool.acquire() as conn: + # Get team information + teams = await conn.fetch( + """ + SELECT id, team_type, settings FROM teams + WHERE id IN ($1, $2) + """, + source_team_id, + target_team_id, + ) + + if len(teams) != 2: + return 0.0 + + source_team = next(t for t in teams if str(t["id"]) == source_team_id) + target_team = next(t for t in teams if str(t["id"]) == target_team_id) + + relevance_factors = [] + + # Factor 1: Team type similarity + type_similarity = ( + 1.0 if source_team["team_type"] == target_team["team_type"] else 0.3 + ) + relevance_factors.append(type_similarity) + + # Factor 2: Knowledge category relevance to target team + target_team_categories = await conn.fetch( + """ + SELECT knowledge_category, COUNT(*) as usage_count + FROM team_knowledge_base + WHERE team_id = $1 + GROUP BY knowledge_category + ORDER BY usage_count DESC + LIMIT 5 + """, + target_team_id, + ) + + target_categories = [ + cat["knowledge_category"] for cat in target_team_categories + ] + category_relevance = ( + 1.0 if knowledge["knowledge_category"] in target_categories else 0.4 + ) + relevance_factors.append(category_relevance) + + # Factor 3: Effectiveness score of source knowledge + effectiveness_factor = knowledge["effectiveness_score"] + relevance_factors.append(effectiveness_factor) + + # Factor 4: Adoption rate in source team (indicates broad utility) + adoption_factor = knowledge["agent_adoption_rate"] + relevance_factors.append(adoption_factor) + + # Calculate weighted relevance + weights = [0.2, 0.3, 0.3, 0.2] + relevance = sum(f * w for f, w in zip(relevance_factors, weights)) + + return min(1.0, max(0.0, relevance)) + + def _create_default_propagation_rules(self) -> List[PropagationRule]: + """Create default propagation rules""" + return [ + # Agent to Team rules + PropagationRule( + source_type="agent", + target_type="team", + min_confidence=0.6, + min_success_correlation=0.0, + min_usage_count=1, + knowledge_categories=list(KnowledgeCategory), + auto_approve=True, + propagation_weight=1.0, + ), + # Team to Organization rules + PropagationRule( + source_type="team", + target_type="organization", + min_confidence=0.7, + min_success_correlation=0.5, + min_usage_count=3, + knowledge_categories=list(KnowledgeCategory), + auto_approve=False, + propagation_weight=0.8, + ), + # Cross-team rules + PropagationRule( + source_type="team", + target_type="team", + min_confidence=0.65, + min_success_correlation=0.4, + min_usage_count=2, + knowledge_categories=[ + KnowledgeCategory.DEVELOPMENT, + KnowledgeCategory.TESTING, + KnowledgeCategory.TROUBLESHOOTING, + ], + auto_approve=False, + propagation_weight=0.6, + ), + ] + + def _row_to_propagation_task(self, row) -> PropagationTask: + """Convert database row to PropagationTask object""" + return PropagationTask( + id=str(row["id"]), + source_type=row["source_type"], + source_id=str(row["source_id"]), + target_type=row["target_type"], + target_id=str(row["target_id"]), + knowledge_type=row["knowledge_type"], + knowledge_content_id=( + str(row["knowledge_content_id"]) if row["knowledge_content_id"] else "" + ), + propagation_method=row["propagation_method"], + propagation_trigger=PropagationTrigger(row["propagation_trigger"]), + confidence_score=row["confidence_score"], + propagation_status=PropagationStatus(row["propagation_status"]), + acceptance_status=AcceptanceStatus(row["acceptance_status"]), + metadata=( + json.loads(row["metadata"]) + if isinstance(row["metadata"], str) + else row["metadata"] + ), + created_at=row["propagated_at"], + processed_at=row["processed_at"], + completed_at=row["completed_at"], + ) diff --git a/services/orchestrator/main.py b/services/orchestrator/main.py index ce26700..bb7fbe2 100644 --- a/services/orchestrator/main.py +++ b/services/orchestrator/main.py @@ -1,6013 +1,6013 @@ -import asyncio -import json -import logging -import os -from collections import defaultdict -from contextlib import asynccontextmanager -from datetime import date, datetime -from decimal import Decimal -from typing import Any, Dict, List, Optional - -import jwt -from fastapi import ( - Body, - Depends, - FastAPI, - File, - Form, - HTTPException, - Path, - Query, - UploadFile, - WebSocket, - WebSocketDisconnect, - status, -) -from fastapi.middleware.cors import CORSMiddleware -from fastapi.responses import HTMLResponse, Response -from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer -from pydantic import BaseModel, Field - -from hierarchy_endpoints import router as hierarchy_router - -from .agent_manager import AgentManager -from .container_manager import ContainerConfig, ContainerStatus, container_manager -from .context_service import ContextService -from .database import get_db_connection -from .knowledge_manager import DocumentMetadata, knowledge_manager -from .rag_integration import RAGContext, rag_system -from .sandbox_manager import AgentSandboxManager -from .task_execution_engine import TaskExecutionEngine -from .task_queue import TaskQueue -from .websocket_manager import ( - UpdateType, - WebSocketUpdate, - notify_agent_status_change, - notify_container_status_change, - notify_knowledge_update, - notify_task_progress, - websocket_manager, -) - -logger = logging.getLogger(__name__) - -# --------------------------------------------------------------------------- -# Auth helpers (Track 3) -# --------------------------------------------------------------------------- -_security = HTTPBearer(auto_error=False) -_jwt_secret = os.environ.get("FUZEFRONT_JWT_SECRET", "") - - -def require_auth(credentials: HTTPAuthorizationCredentials = Depends(_security)): - """Verify FuzeFront JWT on mutating endpoints. Disabled when secret not set (dev).""" - if not _jwt_secret: - return None # Auth disabled when secret not configured (dev mode) - if not credentials: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing token" - ) - try: - payload = jwt.decode(credentials.credentials, _jwt_secret, algorithms=["HS256"]) - return payload - except jwt.ExpiredSignatureError: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, detail="Token expired" - ) - except jwt.InvalidTokenError: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token" - ) - - -# --------------------------------------------------------------------------- -# Agent relay state (Track 4) -# --------------------------------------------------------------------------- -# agent_id -> list of subscriber WebSockets watching that agent's session -agent_relay_subscribers: Dict[str, List[WebSocket]] = defaultdict(list) - - -# Pydantic models for API documentation -class AgentCreateRequest(BaseModel): - name: str = Field(..., description="Agent name") - role: str = Field(..., description="Agent role (e.g., 'Senior React Developer')") - type: str = Field(..., description="Agent type (e.g., 'developer', 'executive')") - config: Dict[str, Any] = Field( - default_factory=dict, description="Agent configuration" - ) - repository_settings: Dict[str, Any] = Field( - default_factory=dict, description="Repository settings" - ) - sandbox_settings: Dict[str, Any] = Field( - default_factory=dict, description="Sandbox settings" - ) - - -class TaskCreateRequest(BaseModel): - title: str = Field(..., description="Task title") - description: str = Field(..., description="Task description") - priority: str = Field( - default="medium", description="Task priority (low, medium, high)" - ) - metadata: Dict[str, Any] = Field( - default_factory=dict, description="Additional task metadata" - ) - - -class HumanResponseRequest(BaseModel): - response: str = Field(..., description="Human response to agent question") - - -class FileOperationApprovalRequest(BaseModel): - approved: bool = Field(..., description="Whether to approve the file operations") - reason: Optional[str] = Field( - None, description="Optional reason for approval/rejection" - ) - - -class ClaudeSessionInputRequest(BaseModel): - input: str = Field(..., description="Input to send to Claude SDK session") - - -class CoordinationRequest(BaseModel): - coordination_mode: str = Field( - default="collaborative", - description="Coordination mode (sequential, parallel, hierarchical, collaborative)", - ) - required_agents: Optional[List[str]] = Field( - None, description="Specific agents to include" - ) - required_skills: Optional[List[str]] = Field( - None, description="Required skills for the task" - ) - - -class AgentCommunicationRequest(BaseModel): - message_type: str = Field( - default="notification", - description="Message type (request, response, notification, question)", - ) - content: str = Field(..., description="Message content") - metadata: Dict[str, Any] = Field( - default_factory=dict, description="Additional metadata" - ) - - -class MCPToolRequest(BaseModel): - tool_name: str = Field(..., description="Name of the MCP tool to call") - arguments: Dict[str, Any] = Field( - default_factory=dict, description="Tool arguments" - ) - - -class AgentMCPSetupRequest(BaseModel): - task_id: str = Field(..., description="Task ID for MCP setup") - session_id: Optional[str] = Field(None, description="Optional session ID") - - -class ConversationCreateRequest(BaseModel): - title: str = "New Conversation" - initial_message: Optional[str] = None - context: Optional[Dict[str, Any]] = None - - -class ConversationMessage(BaseModel): - role: str # 'user' or 'agent' - content: str - metadata: Optional[Dict[str, Any]] = None - - -class ChatMessageRequest(BaseModel): - content: str - metadata: Optional[Dict[str, Any]] = None - - -# Model Configuration Models -class ProviderCredentialsRequest(BaseModel): - provider: str = Field( - ..., description="Model provider (anthropic, openai, google, etc.)" - ) - api_key: str = Field(..., description="API key for the provider") - endpoint_url: Optional[str] = Field(None, description="Custom endpoint URL") - additional_config: Dict[str, Any] = Field( - default_factory=dict, description="Additional provider configuration" - ) - - -class AgentModelConfigRequest(BaseModel): - primary_model: str = Field(..., description="Primary model ID") - fallback_models: List[str] = Field( - default_factory=list, description="Fallback model IDs" - ) - temperature: float = Field( - default=0.7, ge=0.0, le=2.0, description="Model temperature" - ) - max_tokens: int = Field( - default=4096, ge=1, le=200000, description="Maximum output tokens" - ) - top_p: float = Field(default=1.0, ge=0.0, le=1.0, description="Top-p sampling") - frequency_penalty: float = Field( - default=0.0, ge=-2.0, le=2.0, description="Frequency penalty" - ) - presence_penalty: float = Field( - default=0.0, ge=-2.0, le=2.0, description="Presence penalty" - ) - custom_instructions: str = Field( - default="", description="Custom instructions for the agent" - ) - use_function_calling: bool = Field( - default=True, description="Enable function calling" - ) - streaming_enabled: bool = Field( - default=True, description="Enable response streaming" - ) - cost_limit_per_task: Optional[float] = Field( - None, ge=0.0, description="Cost limit per task in USD" - ) - - -class TaskCostEstimateRequest(BaseModel): - task_description: str = Field(..., description="Description of the task") - estimated_complexity: str = Field( - default="medium", - description="Estimated complexity (low, medium, high, very_high)", - ) - - -# Response models -class AgentResponse(BaseModel): - agent_id: str - status: str - agent: Dict[str, Any] - - -class TaskResponse(BaseModel): - task_id: str - status: str - - -class CoordinationResponse(BaseModel): - task_id: str - coordination_session_id: Optional[str] = None - status: str - coordination_mode: Optional[str] = None - message: Optional[str] = None - - -# Goals Management API Models -class GoalCreateRequest(BaseModel): - title: str = Field(..., description="Goal title") - description: str = Field(..., description="Goal description") - goal_type: str = Field( - default="business", - description="Goal type (business, technical, growth, operational)", - ) - target_value: Optional[Decimal] = Field( - None, description="Target value (e.g., 100000 for $100K)" - ) - target_unit: Optional[str] = Field( - None, description="Target unit (e.g., 'USD', 'users', '%')" - ) - target_deadline: Optional[date] = Field(None, description="Target completion date") - priority_level: int = Field( - default=5, ge=1, le=10, description="Priority level (1-10)" - ) - success_criteria: Optional[Dict[str, Any]] = Field( - default=None, description="Success criteria" - ) - assigned_teams: Optional[List[str]] = Field( - default=None, description="Assigned team IDs" - ) - goal_owner_agent_id: Optional[str] = Field(None, description="Goal owner agent ID") - stakeholder_agents: Optional[List[str]] = Field( - default=None, description="Stakeholder agent IDs" - ) - tags: Optional[List[str]] = Field(default=None, description="Goal tags") - metadata: Optional[Dict[str, Any]] = Field( - default=None, description="Additional metadata" - ) - - -class GoalUpdateRequest(BaseModel): - progress_percentage: Optional[Decimal] = Field( - None, ge=0, le=100, description="Progress percentage" - ) - current_value: Optional[Decimal] = Field(None, description="Current value") - completion_confidence: Optional[Decimal] = Field( - None, ge=0, le=1, description="Completion confidence" - ) - notes: Optional[str] = Field(None, description="Progress notes") - - -class MilestoneCreateRequest(BaseModel): - title: str = Field(..., description="Milestone title") - description: str = Field(..., description="Milestone description") - target_date: date = Field(..., description="Target completion date") - milestone_type: str = Field(default="deliverable", description="Milestone type") - success_criteria: Optional[Dict[str, Any]] = Field( - default=None, description="Success criteria" - ) - deliverables: Optional[List[Dict[str, Any]]] = Field( - default=None, description="Expected deliverables" - ) - dependencies: Optional[List[Dict[str, Any]]] = Field( - default=None, description="Dependencies" - ) - assigned_teams: Optional[List[str]] = Field( - default=None, description="Assigned teams" - ) - responsible_agent_id: Optional[str] = Field(None, description="Responsible agent") - priority_level: int = Field(default=5, ge=1, le=10, description="Priority level") - weight_in_goal: Optional[Decimal] = Field( - None, ge=0, le=100, description="Weight in goal (%)" - ) - - -class TaskFromMilestoneRequest(BaseModel): - title: str = Field(..., description="Task title") - description: str = Field(..., description="Task description") - task_type: str = Field(default="development", description="Task type") - complexity_level: str = Field(default="medium", description="Complexity level") - estimated_hours: Optional[Decimal] = Field(None, description="Estimated hours") - due_date: Optional[date] = Field(None, description="Due date") - assigned_team_id: Optional[str] = Field(None, description="Assigned team ID") - assigned_agent_id: Optional[str] = Field(None, description="Assigned agent ID") - priority: int = Field(default=5, ge=1, le=10, description="Priority") - requirements: Optional[Dict[str, Any]] = Field( - default=None, description="Requirements" - ) - acceptance_criteria: Optional[List[Dict[str, Any]]] = Field( - default=None, description="Acceptance criteria" - ) - dependencies: Optional[List[Dict[str, Any]]] = Field( - default=None, description="Dependencies" - ) - - -class GoalConversationCreateRequest(BaseModel): - conversation_type: str = Field(default="planning", description="Conversation type") - conversation_title: str = Field(..., description="Conversation title") - initial_context: Optional[Dict[str, Any]] = Field( - default=None, description="Initial context" - ) - participants: Optional[List[Dict[str, Any]]] = Field( - default=None, description="Participants" - ) - - -class ConversationMessageRequest(BaseModel): - message_type: str = Field(default="human", description="Message type") - sender_name: str = Field(..., description="Sender name") - content: str = Field(..., description="Message content") - metadata: Optional[Dict[str, Any]] = Field( - default=None, description="Message metadata" - ) - references: Optional[List[str]] = Field( - default=None, description="Referenced message IDs" - ) - - -class ProgressUpdateRequest(BaseModel): - progress_percentage: Optional[Decimal] = Field( - None, ge=0, le=100, description="Progress percentage" - ) - current_value: Optional[Decimal] = Field(None, description="Current value") - milestone_id: Optional[str] = Field(None, description="Associated milestone ID") - notes: Optional[str] = Field(None, description="Progress notes") - confidence_score: Optional[Decimal] = Field( - None, ge=0, le=1, description="Confidence score" - ) - trigger_alerts: bool = Field(default=True, description="Whether to trigger alerts") - - -@asynccontextmanager -async def lifespan(app: FastAPI): - # Startup - database_url = os.getenv( - "DATABASE_URL", "postgresql://postgres:password@postgres:5432/ai_context" - ) - - app.state.agent_manager = AgentManager(database_url) - app.state.task_queue = TaskQueue() - app.state.context_service = ContextService() - - # Initialize sandbox manager - app.state.sandbox_manager = AgentSandboxManager(database_url) - await app.state.sandbox_manager.start() - - # Start WebSocket manager background cleanup task - await websocket_manager.start() - - # Initialize task execution engine - app.state.task_execution_engine = TaskExecutionEngine(app.state.sandbox_manager) - await app.state.task_execution_engine.start() - - # Initialize multi-agent coordinator - from .multi_agent_coordinator import integrate_multi_agent_coordination - - app.state.multi_agent_coordinator = integrate_multi_agent_coordination( - app.state.task_execution_engine - ) - await app.state.multi_agent_coordinator.start() - - # Initialize knowledge management system - try: - from .context_enhancement_service import ContextEnhancementService - from .knowledge_notification_service import KnowledgeNotificationService - from .knowledge_propagation_engine import KnowledgePropagationEngine - from .organization_rag_manager import OrganizationRAGManager - from .task_knowledge_extractor import TaskKnowledgeExtractor - from .team_knowledge_manager import TeamKnowledgeManager - - app.state.org_rag_manager = OrganizationRAGManager(database_url) - await app.state.org_rag_manager.initialize() - - app.state.team_knowledge_manager = TeamKnowledgeManager(database_url) - await app.state.team_knowledge_manager.initialize() - - app.state.knowledge_propagation_engine = KnowledgePropagationEngine( - database_url, app.state.org_rag_manager, app.state.team_knowledge_manager - ) - await app.state.knowledge_propagation_engine.initialize() - - app.state.notification_service = KnowledgeNotificationService(database_url) - await app.state.notification_service.initialize() - - app.state.task_knowledge_extractor = TaskKnowledgeExtractor( - database_url, - app.state.org_rag_manager, - app.state.team_knowledge_manager, - app.state.knowledge_propagation_engine, - ) - await app.state.task_knowledge_extractor.initialize() - - app.state.context_enhancement_service = ContextEnhancementService( - database_url, app.state.org_rag_manager, app.state.team_knowledge_manager - ) - await app.state.context_enhancement_service.initialize() - - # Initialize knowledge analytics service - from .knowledge_analytics_service import KnowledgeAnalyticsService - - app.state.knowledge_analytics_service = KnowledgeAnalyticsService(database_url) - await app.state.knowledge_analytics_service.initialize() - - logger.info("Knowledge management system initialized successfully") - - except Exception as e: - logger.warning(f"Failed to initialize knowledge management system: {e}") - - # Initialize goals management system - try: - from .goal_conversation_service import GoalConversationService - from .goal_tracking_service import GoalTrackingService - from .goals_management_service import GoalsManagementService - from .milestone_task_engine import MilestoneTaskEngine - - app.state.goals_service = GoalsManagementService(database_url) - await app.state.goals_service.initialize() - - app.state.milestone_task_engine = MilestoneTaskEngine(database_url) - await app.state.milestone_task_engine.initialize() - - app.state.goal_conversation_service = GoalConversationService(database_url) - await app.state.goal_conversation_service.initialize() - - app.state.goal_tracking_service = GoalTrackingService(database_url) - await app.state.goal_tracking_service.initialize() - - logger.info("Goals management system initialized successfully") - - except Exception as e: - logger.warning(f"Failed to initialize goals management system: {e}") - - # Connect components - app.state.task_queue.set_task_execution_engine(app.state.task_execution_engine) - await app.state.agent_manager.set_sandbox_manager(app.state.sandbox_manager) - - # Initialize IzzyAI CEO on startup - try: - await app.state.agent_manager.create_agent( - name="IzzyAI", - role="Digital CEO", - type="executive", - config={ - "model": "claude-sonnet-4-20250514", - "temperature": 0.7, - "tools": [ - "strategic_planning", - "resource_allocation", - "team_management", - ], - }, - ) - except Exception as e: - print(f"Warning: Could not create IzzyAI CEO: {e}") - - yield - - # Shutdown - await app.state.multi_agent_coordinator.stop() - await app.state.task_execution_engine.stop() - await app.state.sandbox_manager.stop() - await app.state.agent_manager.shutdown_all() - await app.state.task_queue.close() - - # Shutdown knowledge management services - try: - if hasattr(app.state, "knowledge_analytics_service"): - await app.state.knowledge_analytics_service.close() - if hasattr(app.state, "context_enhancement_service"): - await app.state.context_enhancement_service.close() - if hasattr(app.state, "task_knowledge_extractor"): - await app.state.task_knowledge_extractor.close() - if hasattr(app.state, "notification_service"): - await app.state.notification_service.close() - if hasattr(app.state, "knowledge_propagation_engine"): - await app.state.knowledge_propagation_engine.close() - if hasattr(app.state, "team_knowledge_manager"): - await app.state.team_knowledge_manager.close() - if hasattr(app.state, "org_rag_manager"): - await app.state.org_rag_manager.close() - logger.info("Knowledge management system shutdown complete") - except Exception as e: - logger.error(f"Error shutting down knowledge management system: {e}") - - # Shutdown goals management services - try: - if hasattr(app.state, "goal_tracking_service"): - await app.state.goal_tracking_service.close() - if hasattr(app.state, "goal_conversation_service"): - await app.state.goal_conversation_service.close() - if hasattr(app.state, "milestone_task_engine"): - await app.state.milestone_task_engine.close() - if hasattr(app.state, "goals_service"): - await app.state.goals_service.close() - logger.info("Goals management system shutdown complete") - except Exception as e: - logger.error(f"Error shutting down goals management system: {e}") - - -app = FastAPI( - title="FuzeAgent Orchestrator API", - description=""" - ## FuzeAgent AI Team Orchestration Platform - - A comprehensive platform for autonomous AI development teams that enables: - - ### 🤖 Autonomous Agent Execution - - **Claude SDK Integration**: Interactive AI development with real-time conversation streaming - - **File Operations Engine**: Safe code changes with human approval workflows - - **Multi-Agent Coordination**: Complex task decomposition and agent collaboration - - ### 🏗️ Core Features - - **Agent Management**: Create, configure, and manage AI development agents - - **Task Orchestration**: Assign and monitor complex development tasks - - **Real-time Monitoring**: WebSocket streaming for live progress updates - - **Human-in-the-Loop**: Seamless approval workflows for critical decisions - - ### 🔗 Integration Capabilities - - **MCP (Model Context Protocol)**: Organizational context for AI agents - - **Git Workflow Management**: Automated repository operations - - **Sandbox Environments**: Isolated development containers - - **Database Integration**: PostgreSQL for persistent storage - - ### 📡 API Categories - - **Agent Management**: Create and manage AI agents - - **Task Execution**: Autonomous task processing - - **File Operations**: Code change management - - **Multi-Agent Coordination**: Team collaboration - - **Real-time Communication**: WebSocket endpoints - - **MCP Integration**: Organizational context tools - - **Goals Management**: Organizational goals, milestones, and task planning - - **Knowledge Management**: RAG system and organizational memory - - **Version**: 2.0.0 (Autonomous Execution) - """, - version="2.0.0", - lifespan=lifespan, - docs_url="/docs", - redoc_url="/redoc", - openapi_tags=[ - {"name": "health", "description": "Health check and system status endpoints"}, - { - "name": "agents", - "description": "AI agent creation, management, and status monitoring", - }, - {"name": "tasks", "description": "Task assignment, execution, and monitoring"}, - { - "name": "autonomous-execution", - "description": "Autonomous task execution with Claude SDK integration", - }, - { - "name": "file-operations", - "description": "File system operations and code change management", - }, - { - "name": "multi-agent-coordination", - "description": "Multi-agent collaboration and task coordination", - }, - { - "name": "real-time", - "description": "WebSocket endpoints for real-time updates", - }, - { - "name": "human-in-loop", - "description": "Human approval workflows and interaction handling", - }, - { - "name": "mcp-integration", - "description": "Model Context Protocol tools and resources", - }, - {"name": "sandboxes", "description": "Sandbox environment management"}, - {"name": "context", "description": "Agent memory and context management"}, - { - "name": "model-configuration", - "description": "AI model configuration and API key management", - }, - { - "name": "knowledge-management", - "description": "Hierarchical knowledge management, RAG, and intelligent notifications", - }, - ], -) - -app.add_middleware( - CORSMiddleware, - allow_origins=[ - "http://localhost:3000", - "http://localhost:3031", - "http://localhost", - ], - allow_credentials=True, - allow_methods=["*"], - allow_headers=["*"], -) - -# Include hierarchy router for organizational visualization -app.include_router(hierarchy_router) - - -# Health check endpoint -@app.get( - "/health", - tags=["health"], - summary="Health Check", - description="Check the health status of the FuzeAgent orchestrator service", - response_description="Service health status", -) -async def health_check(): - """ - Health check endpoint that returns the current status of the orchestrator service. - - Returns: - dict: Service health status and basic information - """ - return { - "status": "healthy", - "service": "orchestrator", - "version": "2.0.0", - "features": { - "autonomous_execution": True, - "multi_agent_coordination": True, - "file_operations": True, - "mcp_integration": True, - "real_time_streaming": True, - }, - # Whether GET /openapi.yaml can answer. An image built without its - # contract is DEGRADED, not dead — the probe still passes (no restart - # can conjure a file the image lacks) but the condition is visible to - # anything that looks, instead of surfacing only as a 503 later. - "openapi": "loaded" if _openapi_document() is not None else "unavailable", - } - - -# --------------------------------------------------------------------------- -# The contract, SERVED. -# -# contracts/openapi.yaml describes this orchestrator's real HTTP surface, with -# the curated descriptions and the irreversibility guidance that -# mcp/tools.overrides.yaml narrows. Committing it is not the same as publishing -# it: consumers — the MCP gateway among them — discover the surface over HTTP. -# -# This is NOT /openapi.json. FastAPI generates that from the code at import -# time; it is accurate about shapes and says nothing about which operations -# dispatch an agent that cannot be recalled. Both are served. This one is the -# contract. -# -# The document is read from the IMAGE, never from a mount, so what this endpoint -# publishes is always the contract this build was compiled against. -# --------------------------------------------------------------------------- -_ORCH_DIR = os.path.dirname(os.path.abspath(__file__)) -_OPENAPI_CANDIDATES = [ - p - for p in [ - os.getenv("OPENAPI_SPEC_PATH"), - os.path.join(_ORCH_DIR, "contracts", "openapi.yaml"), - os.path.join(_ORCH_DIR, "..", "..", "contracts", "openapi.yaml"), - ] - if p -] - - -def _openapi_document(): - """Return the OpenAPI document text, or None when the image lacks it.""" - for path in _OPENAPI_CANDIDATES: - try: - with open(path, "r", encoding="utf-8") as fh: - return fh.read() - except OSError: - continue - return None - - -@app.get( - "/openapi.yaml", - tags=["health"], - summary="This OpenAPI Document", - description=( - "Serve contracts/openapi.yaml — the curated contract, as distinct from " - "FastAPI's auto-generated /openapi.json." - ), - include_in_schema=False, -) -async def get_openapi_document(): - doc = _openapi_document() - if doc is None: - logger.error("OpenAPI document not found; tried %s", _OPENAPI_CANDIDATES) - # 503, not 500 and not a crash: the service is otherwise functional and - # no restart can produce a spec the image does not contain. - raise HTTPException( - status_code=503, - detail=( - "openapi_document_unavailable: this image was built without " - "contracts/openapi.yaml. Rebuild with the repo root as the Docker " - "context so the contract is copied in." - ), - ) - return Response(content=doc, media_type="application/yaml") - - -# WebSocket for real-time updates -@app.websocket("/ws") -async def websocket_endpoint(websocket: WebSocket): - await websocket.accept() - try: - while True: - # Send agent updates to UI - updates = await app.state.agent_manager.get_updates() - await websocket.send_json(updates) - await asyncio.sleep(1) - except Exception as e: - print(f"WebSocket error: {e}") - finally: - await websocket.close() - - -# WebSocket for task execution updates -@app.websocket("/ws/tasks/{task_id}") -async def task_websocket_endpoint(websocket: WebSocket, task_id: str): - """WebSocket endpoint for real-time task execution updates""" - await websocket.accept() - try: - while True: - # Get task execution status - try: - status = await app.state.task_queue.get_execution_status(task_id) - await websocket.send_json( - {"type": "status_update", "task_id": task_id, "data": status} - ) - - # If task is completed or failed, send final update and close - if status.get("status") in ["completed", "failed", "cancelled"]: - await websocket.send_json( - { - "type": "task_finished", - "task_id": task_id, - "final_status": status.get("status"), - } - ) - break - - except Exception as e: - await websocket.send_json( - {"type": "error", "task_id": task_id, "error": str(e)} - ) - - await asyncio.sleep(2) # Update every 2 seconds - - except Exception as e: - print(f"Task WebSocket error for {task_id}: {e}") - finally: - await websocket.close() - - -# WebSocket for real-time Claude SDK conversation streaming -@app.websocket("/ws/tasks/{task_id}/conversation") -async def conversation_websocket_endpoint(websocket: WebSocket, task_id: str): - """WebSocket endpoint for real-time Claude SDK conversation streaming""" - await websocket.accept() - try: - # Get execution context - execution = app.state.task_execution_engine.active_executions.get(task_id) - if not execution: - await websocket.send_json( - {"type": "error", "message": f"Task {task_id} not found or not active"} - ) - await websocket.close() - return - - # Wait for Claude SDK session to be available - while not execution.claude_session_id and execution.status not in [ - "completed", - "failed", - "cancelled", - ]: - await asyncio.sleep(1) - - if not execution.claude_session_id: - await websocket.send_json( - { - "type": "error", - "message": "No active Claude SDK session for this task", - } - ) - await websocket.close() - return - - # Stream Claude SDK output - claude_sdk_manager = execution.claude_sdk_manager - if claude_sdk_manager: - try: - async for output_chunk in claude_sdk_manager.stream_output( - execution.claude_session_id - ): - await websocket.send_json( - { - "type": "claude_output", - "task_id": task_id, - "content": output_chunk, - "timestamp": datetime.now().isoformat(), - } - ) - - # Session ended - await websocket.send_json( - { - "type": "conversation_ended", - "task_id": task_id, - "timestamp": datetime.now().isoformat(), - } - ) - - except Exception as e: - await websocket.send_json( - { - "type": "error", - "message": f"Error streaming conversation: {str(e)}", - } - ) - - except Exception as e: - print(f"Conversation WebSocket error for {task_id}: {e}") - finally: - await websocket.close() - - -# WebSocket for file operations streaming -@app.websocket("/ws/tasks/{task_id}/file-operations") -async def file_operations_websocket_endpoint(websocket: WebSocket, task_id: str): - """WebSocket endpoint for real-time file operations updates""" - await websocket.accept() - try: - # Get execution context - execution = app.state.task_execution_engine.active_executions.get(task_id) - if not execution: - await websocket.send_json( - {"type": "error", "message": f"Task {task_id} not found or not active"} - ) - await websocket.close() - return - - file_ops_engine = execution.file_operations_engine - if not file_ops_engine: - await websocket.send_json( - { - "type": "error", - "message": "No file operations engine available for this task", - } - ) - await websocket.close() - return - - last_batch_count = 0 - - while execution.status not in ["completed", "failed", "cancelled"]: - try: - # Get pending operations - pending_operations = file_ops_engine.get_pending_operations() - applied_operations = file_ops_engine.get_applied_operations() - - current_batch_count = len(pending_operations) + len(applied_operations) - - # Send updates if there are new operations - if current_batch_count > last_batch_count: - # Send pending operations - for batch in pending_operations: - # Get diff preview - diffs = await file_ops_engine.get_file_diff_preview( - batch.batch_id - ) - - await websocket.send_json( - { - "type": "pending_operations", - "task_id": task_id, - "batch_id": batch.batch_id, - "description": batch.description, - "requires_approval": batch.requires_approval, - "operations_count": len(batch.operations), - "file_diffs": diffs, - "timestamp": batch.created_at.isoformat(), - } - ) - - # Send applied operations - for batch in applied_operations: - await websocket.send_json( - { - "type": "applied_operations", - "task_id": task_id, - "batch_id": batch.batch_id, - "description": batch.description, - "operations_count": len(batch.operations), - "applied_at": ( - batch.applied_at.isoformat() - if batch.applied_at - else None - ), - "timestamp": batch.created_at.isoformat(), - } - ) - - last_batch_count = current_batch_count - - await asyncio.sleep(1) # Check every second - - except Exception as e: - await websocket.send_json( - { - "type": "error", - "message": f"Error getting file operations: {str(e)}", - } - ) - - # Task completed - await websocket.send_json( - { - "type": "task_completed", - "task_id": task_id, - "final_status": execution.status.value, - "timestamp": datetime.now().isoformat(), - } - ) - - except Exception as e: - print(f"File operations WebSocket error for {task_id}: {e}") - finally: - await websocket.close() - - -# Agent Management Endpoints -@app.post( - "/agents", - tags=["agents"], - summary="Create AI Agent", - description="Create a new AI agent with repository and sandbox settings", - response_model=AgentResponse, -) -async def create_agent(agent_config: AgentCreateRequest, _auth=Depends(require_auth)): - """Create a new AI agent with repository and sandbox settings""" - try: - agent = await app.state.agent_manager.create_agent(**agent_config) - return { - "agent_id": agent.id, - "status": "created", - "agent": { - "id": agent.id, - "name": agent_config.get("name"), - "role": agent_config.get("role"), - "type": agent_config.get("type"), - "repository_settings": agent_config.get("repository_settings", {}), - "sandbox_settings": agent_config.get("sandbox_settings", {}), - "created_at": ( - agent.created_at if hasattr(agent, "created_at") else None - ), - }, - } - except Exception as e: - raise HTTPException(status_code=400, detail=f"Failed to create agent: {str(e)}") - - -@app.get( - "/agents", - tags=["agents"], - summary="List All Agents", - description="Get a list of all AI agents and their current status", -) -async def list_agents(): - """List all agents and their status""" - return await app.state.agent_manager.list_agents() - - -@app.post( - "/agents/{agent_id}/tasks", - tags=["tasks"], - summary="Assign Task to Agent", - description="Assign a specific task to an AI agent", - response_model=TaskResponse, -) -async def assign_task( - agent_id: str = Path(..., description="Agent ID"), - task: TaskCreateRequest = Body(...), - _auth=Depends(require_auth), -): - """Assign a task to an agent""" - task_id = await app.state.task_queue.assign_task(agent_id, task) - return {"task_id": task_id, "status": "assigned"} - - -@app.get("/agents/{agent_id}/status") -async def get_agent_status(agent_id: str): - """Get detailed agent status""" - return await app.state.agent_manager.get_agent_status(agent_id) - - -@app.get("/agents/{agent_id}/tasks") -async def get_agent_tasks(agent_id: str): - """Get tasks assigned to an agent""" - try: - # This would normally query the database for tasks assigned to the agent - # For now, return mock data - return [ - { - "id": "1", - "title": "Strategic Planning Q4 2025", - "description": "Develop comprehensive strategic plan for Q4 2025 expansion", - "status": "completed", - "priority": "high", - "created_at": "2025-08-05T09:00:00Z", - "completed_at": "2025-08-05T17:30:00Z", - }, - { - "id": "2", - "title": "Team Performance Review", - "description": "Conduct quarterly performance review for all team leads", - "status": "in_progress", - "priority": "medium", - "created_at": "2025-08-06T08:00:00Z", - }, - ] - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get agent tasks: {str(e)}" - ) - - -@app.get("/teams") -async def list_teams(): - """List all teams""" - try: - # This would normally query the database for teams - # For now, return mock data - return [ - { - "id": "1", - "name": "Executive Team", - "description": "Strategic leadership and decision making", - }, - { - "id": "2", - "name": "Development Team", - "description": "Frontend, backend, and full-stack development", - }, - { - "id": "3", - "name": "Quality Assurance", - "description": "Testing, quality control, and bug detection", - }, - { - "id": "4", - "name": "DevOps Team", - "description": "Infrastructure, deployment, and system operations", - }, - { - "id": "5", - "name": "Business Team", - "description": "Marketing, sales, and customer relations", - }, - ] - except Exception as e: - raise HTTPException(status_code=500, detail=f"Failed to list teams: {str(e)}") - - -@app.get("/agent-templates") -async def list_agent_templates(): - """List available agent templates""" - try: - return [ - { - "id": "react_developer", - "name": "React Developer", - "description": "Frontend developer specialized in React and TypeScript", - "type": "developer", - "defaultConfig": { - "model": "claude-sonnet-4-20250514", - "temperature": 0.7, - "tools": ["code_generation", "code_review", "debugging", "testing"], - "goal": "Build responsive and performant React applications", - "backstory": "Experienced frontend developer with deep knowledge of React ecosystem", - }, - }, - { - "id": "python_developer", - "name": "Python Developer", - "description": "Backend developer specialized in Python and FastAPI", - "type": "developer", - "defaultConfig": { - "model": "claude-sonnet-4-20250514", - "temperature": 0.7, - "tools": [ - "code_generation", - "api_development", - "database_design", - "testing", - ], - "goal": "Develop robust and scalable backend systems", - "backstory": "Senior Python developer with expertise in FastAPI and databases", - }, - }, - { - "id": "qa_engineer", - "name": "QA Engineer", - "description": "Quality assurance engineer focused on testing and automation", - "type": "qa", - "defaultConfig": { - "model": "claude-sonnet-4-20250514", - "temperature": 0.6, - "tools": [ - "test_automation", - "bug_reporting", - "quality_analysis", - "performance_testing", - ], - "goal": "Ensure high quality and reliability of software products", - "backstory": "Experienced QA engineer with expertise in automated testing frameworks", - }, - }, - { - "id": "devops_engineer", - "name": "DevOps Engineer", - "description": "Infrastructure and deployment specialist", - "type": "devops", - "defaultConfig": { - "model": "claude-sonnet-4-20250514", - "temperature": 0.5, - "tools": [ - "infrastructure_management", - "deployment", - "monitoring", - "security", - ], - "goal": "Maintain reliable and scalable infrastructure", - "backstory": "DevOps engineer with expertise in cloud platforms and CI/CD", - }, - }, - ] - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to list agent templates: {str(e)}" - ) - - -@app.get("/tasks") -async def list_tasks(): - """List all tasks""" - return await app.state.task_queue.list_tasks() - - -@app.get("/tasks/{task_id}") -async def get_task(task_id: str): - """Get task details""" - return await app.state.task_queue.get_task(task_id) - - -# Autonomous Execution Endpoints -@app.post("/agents/from-template") -async def create_agent_from_template(request: dict): - """Create agent from template with repository settings""" - try: - # Extract template data - template_id = request.get("template_id") - name = request.get("name") - team_id = request.get("team_id") - overrides = request.get("overrides", {}) - - # Get template configuration - template_config = await app.state.agent_manager.get_template_config(template_id) - if not template_config: - raise HTTPException( - status_code=404, detail=f"Template {template_id} not found" - ) - - # Build agent configuration - agent_config = { - "name": name, - "role": template_config.get("role", template_id.replace("_", " ").title()), - "type": template_config.get("type", "specialized"), - "template_id": template_id, - "team_id": team_id, - "config": {**template_config.get("config", {}), **overrides}, - "repository_settings": request.get("repository_settings", {}), - "sandbox_settings": { - "base_image": f"fuzeagent/dev-{template_id.split('_')[0]}:latest", - "resource_limits": template_config.get( - "resource_limits", {"memory": "2Gi", "cpu": "1.0", "disk": "10Gi"} - ), - "auto_cleanup": "24h", - }, - } - - # Create agent - agent = await app.state.agent_manager.create_agent(**agent_config) - - return { - "agent_id": agent.id, - "status": "created", - "agent": agent_config, - "template_id": template_id, - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to create agent from template: {str(e)}" - ) - - -@app.get("/templates") -async def get_agent_templates(): - """Get available agent templates""" - return await app.state.agent_manager.get_available_templates() - - -@app.post( - "/tasks/{task_id}/execute", - tags=["autonomous-execution"], - summary="Start Autonomous Task Execution", - description="Begin autonomous execution of a task using Claude SDK integration", - response_model=TaskResponse, -) -async def start_task_execution( - task_id: str = Path(..., description="Task ID to execute") -): - """Start autonomous execution of a task""" - try: - # This will be handled by the TaskExecutionEngine - result = await app.state.task_queue.start_autonomous_execution(task_id) - return {"task_id": task_id, "status": "execution_started", "result": result} - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to start task execution: {str(e)}" - ) - - -@app.get("/tasks/{task_id}/status") -async def get_task_execution_status(task_id: str): - """Get detailed task execution status""" - try: - status = await app.state.task_queue.get_execution_status(task_id) - return status - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get task status: {str(e)}" - ) - - -@app.get("/tasks/{task_id}/iterations") -async def get_task_iterations(task_id: str): - """Get task iteration history""" - try: - iterations = await app.state.task_queue.get_task_iterations(task_id) - return {"task_id": task_id, "iterations": iterations} - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get task iterations: {str(e)}" - ) - - -@app.get("/agents/{agent_id}/sandbox") -async def get_agent_sandbox(agent_id: str): - """Get agent sandbox information""" - try: - sandbox_info = await app.state.agent_manager.get_agent_sandbox(agent_id) - return sandbox_info - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get agent sandbox: {str(e)}" - ) - - -# Additional endpoints for UI support -@app.put("/tasks/{task_id}") -async def update_task(task_id: str, update_data: dict): - """Update task status and result""" - await app.state.task_queue.update_task_status( - task_id=task_id, - status=update_data.get("status"), - result=update_data.get("result"), - ) - return {"status": "updated"} - - -@app.post("/context/interactions") -async def store_interaction(interaction_data: dict): - """Store agent interaction""" - interaction_id = await app.state.context_service.store_interaction( - agent_id=interaction_data.get("agent_id"), - content=interaction_data.get("content"), - metadata=interaction_data.get("metadata", {}), - ) - return {"interaction_id": interaction_id} - - -@app.get("/context") -async def get_context(query: str, agent_id: str = None): - """Get relevant context for a query""" - context = await app.state.context_service.get_context(query, agent_id) - return context - - -@app.get("/agents/{agent_id}/memory") -async def get_agent_memory(agent_id: str, limit: int = 10): - """Get agent memory""" - memory = await app.state.context_service.get_agent_memory(agent_id, limit) - return memory - - -# Agent Conversation Endpoints -@app.get( - "/agents/{agent_id}/conversations", - tags=["conversations"], - summary="Get Agent Conversations", - description="Get all conversations for a specific agent", -) -async def get_agent_conversations(agent_id: str): - """Get all conversations for a specific agent""" - try: - async with get_db_connection() as conn: - conversations = await conn.fetch( - """ - SELECT cs.*, COUNT(ac.id) as message_count, - (SELECT content FROM agent_conversations - WHERE session_id = cs.id - ORDER BY created_at DESC LIMIT 1) as last_message - FROM chat_sessions cs - LEFT JOIN agent_conversations ac ON cs.id = ac.session_id - WHERE cs.agent_id = $1 - GROUP BY cs.id - ORDER BY cs.last_activity DESC - """, - agent_id, - ) - - return [dict(row) for row in conversations] - - except Exception as e: - logger.error(f"Error getting agent conversations: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/agents/{agent_id}/conversations", - tags=["conversations"], - summary="Create New Agent Conversation", - description="Create a new conversation with an agent", -) -async def create_agent_conversation(agent_id: str, request: ConversationCreateRequest): - """Create a new conversation with an agent""" - try: - async with get_db_connection() as conn: - # Create new chat session - session_id = await conn.fetchval( - """ - INSERT INTO chat_sessions (agent_id, session_name, context, status) - VALUES ($1, $2, $3, 'active') - RETURNING id - """, - agent_id, - request.title, - request.context or {}, - ) - - # Add initial message if provided - if request.initial_message: - await conn.execute( - """ - INSERT INTO agent_conversations (session_id, agent_id, message_type, content) - VALUES ($1, $2, 'system', $3) - """, - session_id, - agent_id, - request.initial_message, - ) - - # Get the created conversation - conversation = await conn.fetchrow( - """ - SELECT * FROM chat_sessions WHERE id = $1 - """, - session_id, - ) - - return dict(conversation) - - except Exception as e: - logger.error(f"Error creating agent conversation: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/agents/{agent_id}/conversations/{conversation_id}/messages", - tags=["conversations"], - summary="Get Conversation Messages", - description="Get all messages in a conversation", -) -async def get_conversation_messages(agent_id: str, conversation_id: str): - """Get all messages in a conversation""" - try: - async with get_db_connection() as conn: - messages = await conn.fetch( - """ - SELECT * FROM agent_conversations - WHERE session_id = $1 AND agent_id = $2 - ORDER BY created_at ASC - """, - conversation_id, - agent_id, - ) - - return [dict(row) for row in messages] - - except Exception as e: - logger.error(f"Error getting conversation messages: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/agents/{agent_id}/conversations/{conversation_id}/messages", - tags=["conversations"], - summary="Send Message to Agent", - description="Send a message to an agent in a conversation", -) -async def send_message_to_agent( - agent_id: str, conversation_id: str, request: ChatMessageRequest -): - """Send a message to an agent in a conversation""" - try: - async with get_db_connection() as conn: - # Insert user message - user_message_id = await conn.fetchval( - """ - INSERT INTO agent_conversations (session_id, agent_id, message_type, content, metadata) - VALUES ($1, $2, 'user', $3, $4) - RETURNING id - """, - conversation_id, - agent_id, - request.content, - request.metadata or {}, - ) - - # Update session last activity - await conn.execute( - """ - UPDATE chat_sessions - SET last_activity = CURRENT_TIMESTAMP, message_count = message_count + 1 - WHERE id = $1 - """, - conversation_id, - ) - - # TODO: Here we would integrate with the actual agent to generate a response - # For now, return a simple acknowledgment - - return { - "id": str(user_message_id), - "status": "sent", - "message": "Message sent to agent", - } - - except Exception as e: - logger.error(f"Error sending message to agent: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get("/knowledge/search") -async def search_knowledge(query: str, limit: int = 10): - """Search knowledge base""" - results = await app.state.context_service.search_knowledge(query, limit) - return results - - -# Human-in-the-loop endpoints -@app.post( - "/tasks/{task_id}/human-response", - tags=["human-in-loop"], - summary="Submit Human Response", - description="Submit human response to a task question or approval request", -) -async def submit_human_response( - task_id: str = Path(..., description="Task ID"), - response_data: HumanResponseRequest = Body(...), -): - """Submit human response to a task question""" - try: - response = response_data.get("response", "") - if not response: - raise HTTPException(status_code=400, detail="Response cannot be empty") - - success = await app.state.task_queue.handle_human_response(task_id, response) - - if success: - return {"status": "success", "message": "Human response submitted"} - else: - raise HTTPException( - status_code=404, - detail="Task not found or not waiting for human response", - ) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to submit human response: {str(e)}" - ) - - -@app.post("/tasks/{task_id}/cancel") -async def cancel_task_execution(task_id: str): - """Cancel autonomous execution of a task""" - try: - success = await app.state.task_queue.cancel_task_execution(task_id) - - if success: - return {"status": "cancelled", "message": "Task execution cancelled"} - else: - raise HTTPException(status_code=404, detail="Task not found or not running") - - except Exception as e: - raise HTTPException(status_code=500, detail=f"Failed to cancel task: {str(e)}") - - -@app.get("/tasks/{task_id}/messages") -async def get_task_messages(task_id: str): - """Get task messages and chat history""" - try: - # This would integrate with the HumanInTheLoopHandler when implemented - # For now, return iteration history which includes human interactions - iterations = await app.state.task_queue.get_task_iterations(task_id) - - messages = [] - for iteration in iterations: - if iteration.get("human_question"): - messages.append( - { - "type": "agent_question", - "content": iteration["human_question"], - "timestamp": iteration["started_at"], - "iteration": iteration["iteration_number"], - } - ) - - if iteration.get("human_response"): - messages.append( - { - "type": "human_response", - "content": iteration["human_response"], - "timestamp": iteration["completed_at"] - or iteration["started_at"], - "iteration": iteration["iteration_number"], - } - ) - - return {"task_id": task_id, "messages": messages} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get task messages: {str(e)}" - ) - - -# Sandbox management endpoints -@app.get("/sandboxes") -async def list_sandboxes(agent_id: str = None, status: str = None): - """List active sandboxes""" - try: - from .sandbox_manager import SandboxStatus - - sandbox_status = None - if status: - try: - sandbox_status = SandboxStatus(status) - except ValueError: - raise HTTPException(status_code=400, detail=f"Invalid status: {status}") - - sandboxes = await app.state.sandbox_manager.list_sandboxes( - agent_id=agent_id, status=sandbox_status - ) - - return { - "sandboxes": [ - { - "sandbox_id": s.sandbox_id, - "agent_id": s.agent_id, - "task_id": s.task_id, - "status": s.status.value, - "workspace_path": s.workspace_path, - "created_at": s.created_at.isoformat(), - "resource_limits": s.resource_limits, - } - for s in sandboxes - ] - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to list sandboxes: {str(e)}" - ) - - -@app.post("/sandboxes/{sandbox_id}/execute") -async def execute_command_in_sandbox(sandbox_id: str, command_data: dict): - """Execute a command in a sandbox""" - try: - command = command_data.get("command") - working_dir = command_data.get("working_dir") - - if not command: - raise HTTPException(status_code=400, detail="Command is required") - - result = await app.state.sandbox_manager.execute_command( - sandbox_id=sandbox_id, command=command, working_dir=working_dir - ) - - return result - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to execute command: {str(e)}" - ) - - -@app.delete("/sandboxes/{sandbox_id}") -async def destroy_sandbox(sandbox_id: str): - """Destroy a sandbox""" - try: - await app.state.sandbox_manager.destroy_sandbox(sandbox_id) - return {"status": "destroyed", "sandbox_id": sandbox_id} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to destroy sandbox: {str(e)}" - ) - - -# Agent registration and communication endpoints -@app.post("/agents/{agent_id}/register") -async def register_agent(agent_id: str, registration_data: dict): - """Register an agent running in a sandbox container""" - try: - # Store agent registration info - # This would typically update the agent's status and capabilities - return { - "status": "registered", - "agent_id": agent_id, - "registered_at": datetime.now().isoformat(), - } - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to register agent: {str(e)}" - ) - - -@app.get("/agents/{agent_id}/next-task") -async def get_next_task_for_agent(agent_id: str): - """Get the next task for an agent to execute""" - try: - # Find pending tasks assigned to this agent - tasks = await app.state.task_queue.get_agent_tasks(agent_id) - pending_tasks = [t for t in tasks if t.get("status") == "pending"] - - if pending_tasks: - # Return the first pending task - task = pending_tasks[0] - # Update status to 'assigned' to prevent double assignment - await app.state.task_queue.update_task_status(task["id"], "assigned") - return task - else: - # No tasks available - return None, 204 - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get next task: {str(e)}" - ) - - -@app.post("/agents/{agent_id}/error") -async def report_agent_error(agent_id: str, error_data: dict): - """Report an error from an agent""" - try: - # Log the error and update agent status - logger.error(f"Agent {agent_id} reported error: {error_data.get('error')}") - - # You might want to store this in a database or alerting system - return {"status": "error_logged", "agent_id": agent_id} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to log agent error: {str(e)}" - ) - - -# Conversation management endpoints -@app.get("/tasks/{task_id}/conversation") -async def get_task_conversation(task_id: str, iteration: int = None, limit: int = 100): - """Get conversation history for a task""" - try: - conversation_history = await app.state.task_execution_engine.conversation_manager.get_conversation_history( - task_id=task_id, iteration_number=iteration, limit=limit - ) - - return {"task_id": task_id, "conversation": conversation_history} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get conversation: {str(e)}" - ) - - -@app.get("/tasks/{task_id}/conversation/summary") -async def get_conversation_summary(task_id: str): - """Get conversation summary with statistics""" - try: - summary = await app.state.task_execution_engine.conversation_manager.get_conversation_summary( - task_id - ) - return summary - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get conversation summary: {str(e)}" - ) - - -@app.get("/tasks/{task_id}/code-generations") -async def get_task_code_generations( - task_id: str, iteration: int = None, file_type: str = None -): - """Get code generations for a task""" - try: - code_generations = await app.state.task_execution_engine.conversation_manager.get_code_generations( - task_id=task_id, iteration_number=iteration, file_type=file_type - ) - - return {"task_id": task_id, "code_generations": code_generations} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get code generations: {str(e)}" - ) - - -@app.get("/agents/{agent_id}/performance") -async def get_agent_performance(agent_id: str, hours: int = 24): - """Get agent performance metrics""" - try: - metrics = await app.state.task_execution_engine.conversation_manager.get_agent_performance_metrics( - agent_id=agent_id, time_range_hours=hours - ) - - return {"agent_id": agent_id, "time_range_hours": hours, "metrics": metrics} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get agent performance: {str(e)}" - ) - - -# File Operations Endpoints -@app.get( - "/tasks/{task_id}/file-operations", - tags=["file-operations"], - summary="Get File Operations", - description="Get file operations for a task with optional status filtering", -) -async def get_task_file_operations( - task_id: str = Path(..., description="Task ID"), - status: Optional[str] = Query( - None, description="Filter by status (pending, applied)" - ), -): - """Get file operations for a task""" - try: - execution = app.state.task_execution_engine.active_executions.get(task_id) - if not execution or not execution.file_operations_engine: - raise HTTPException( - status_code=404, detail="Task not found or no file operations available" - ) - - file_ops_engine = execution.file_operations_engine - - if status == "pending": - operations = file_ops_engine.get_pending_operations() - elif status == "applied": - operations = file_ops_engine.get_applied_operations() - else: - # Get all operations - pending = file_ops_engine.get_pending_operations() - applied = file_ops_engine.get_applied_operations() - operations = pending + applied - - # Convert to dict format - operations_data = [] - for batch in operations: - operations_data.append( - { - "batch_id": batch.batch_id, - "task_id": batch.task_id, - "agent_id": batch.agent_id, - "description": batch.description, - "requires_approval": batch.requires_approval, - "approval_status": batch.approval_status.value, - "operations_count": len(batch.operations), - "created_at": batch.created_at.isoformat(), - "applied_at": ( - batch.applied_at.isoformat() if batch.applied_at else None - ), - } - ) - - return {"task_id": task_id, "operations": operations_data} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get file operations: {str(e)}" - ) - - -@app.get("/tasks/{task_id}/file-operations/{batch_id}/preview") -async def get_file_operations_preview(task_id: str, batch_id: str): - """Get preview of file changes for a batch""" - try: - execution = app.state.task_execution_engine.active_executions.get(task_id) - if not execution or not execution.file_operations_engine: - raise HTTPException( - status_code=404, detail="Task not found or no file operations available" - ) - - file_ops_engine = execution.file_operations_engine - diffs = await file_ops_engine.get_file_diff_preview(batch_id) - - return {"task_id": task_id, "batch_id": batch_id, "file_diffs": diffs} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get file preview: {str(e)}" - ) - - -@app.post( - "/tasks/{task_id}/file-operations/{batch_id}/approve", - tags=["file-operations", "human-in-loop"], - summary="Approve File Operations", - description="Approve or reject file operations from Claude SDK", -) -async def approve_file_operations( - task_id: str = Path(..., description="Task ID"), - batch_id: str = Path(..., description="Batch ID"), - approval_data: FileOperationApprovalRequest = Body(...), -): - """Approve or reject file operations""" - try: - approved = approval_data.get("approved", False) - - execution = app.state.task_execution_engine.active_executions.get(task_id) - if not execution or not execution.file_operations_engine: - raise HTTPException( - status_code=404, detail="Task not found or no file operations available" - ) - - file_ops_engine = execution.file_operations_engine - success = await file_ops_engine.approve_operations(batch_id, approved) - - if success: - # Also notify Claude SDK if there's an active session - if execution.claude_sdk_manager and execution.claude_session_id: - await execution.claude_sdk_manager.approve_file_operations( - execution.claude_session_id, batch_id, approved - ) - - return { - "task_id": task_id, - "batch_id": batch_id, - "approved": approved, - "status": "success", - } - else: - raise HTTPException(status_code=400, detail="Failed to process approval") - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to approve file operations: {str(e)}" - ) - - -@app.post("/tasks/{task_id}/file-operations/{batch_id}/rollback") -async def rollback_file_operations(task_id: str, batch_id: str): - """Rollback applied file operations""" - try: - execution = app.state.task_execution_engine.active_executions.get(task_id) - if not execution or not execution.file_operations_engine: - raise HTTPException( - status_code=404, detail="Task not found or no file operations available" - ) - - file_ops_engine = execution.file_operations_engine - success = await file_ops_engine.rollback_operations(batch_id) - - if success: - return {"task_id": task_id, "batch_id": batch_id, "status": "rolled_back"} - else: - raise HTTPException(status_code=400, detail="Failed to rollback operations") - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to rollback file operations: {str(e)}" - ) - - -# Claude SDK Session Management Endpoints -@app.get("/tasks/{task_id}/claude-session") -async def get_claude_session_status(task_id: str): - """Get Claude SDK session status for a task""" - try: - execution = app.state.task_execution_engine.active_executions.get(task_id) - if ( - not execution - or not execution.claude_sdk_manager - or not execution.claude_session_id - ): - raise HTTPException( - status_code=404, detail="No active Claude SDK session for this task" - ) - - status = await execution.claude_sdk_manager.get_session_status( - execution.claude_session_id - ) - return status - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get Claude session status: {str(e)}" - ) - - -@app.post("/tasks/{task_id}/claude-session/input") -async def send_claude_session_input(task_id: str, input_data: dict): - """Send input to Claude SDK session""" - try: - user_input = input_data.get("input", "") - if not user_input: - raise HTTPException(status_code=400, detail="Input cannot be empty") - - execution = app.state.task_execution_engine.active_executions.get(task_id) - if ( - not execution - or not execution.claude_sdk_manager - or not execution.claude_session_id - ): - raise HTTPException( - status_code=404, detail="No active Claude SDK session for this task" - ) - - success = await execution.claude_sdk_manager.send_input( - execution.claude_session_id, user_input - ) - - if success: - return {"task_id": task_id, "status": "input_sent", "input": user_input} - else: - raise HTTPException( - status_code=400, detail="Failed to send input to Claude session" - ) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to send Claude session input: {str(e)}" - ) - - -# MCP Integration Endpoints -@app.get("/mcp/tools") -async def get_mcp_tools(): - """Get available MCP tools""" - try: - from .mcp_integration import FuzeAgentMCPServer - - mcp_server = FuzeAgentMCPServer() - tools = [ - { - "name": tool.name, - "description": tool.description, - "input_schema": tool.input_schema, - } - for tool in mcp_server.tools - ] - - return {"tools": tools} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get MCP tools: {str(e)}" - ) - - -@app.post( - "/mcp/call-tool", - tags=["mcp-integration"], - summary="Call MCP Tool", - description="Execute an MCP tool to access organizational context", -) -async def call_mcp_tool(tool_request: MCPToolRequest = Body(...)): - """Call an MCP tool""" - try: - from .mcp_integration import FuzeAgentMCPServer - - tool_name = tool_request.get("tool_name") - arguments = tool_request.get("arguments", {}) - - if not tool_name: - raise HTTPException(status_code=400, detail="tool_name is required") - - mcp_server = FuzeAgentMCPServer() - result = await mcp_server.handle_tool_call(tool_name, arguments) - - return result - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to call MCP tool: {str(e)}" - ) - - -@app.get("/mcp/resources") -async def get_mcp_resources(): - """Get available MCP resources""" - try: - from .mcp_integration import FuzeAgentMCPServer - - mcp_server = FuzeAgentMCPServer() - resources = [ - { - "uri": resource.uri, - "name": resource.name, - "description": resource.description, - "mime_type": resource.mime_type, - } - for resource in mcp_server.resources - ] - - return {"resources": resources} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get MCP resources: {str(e)}" - ) - - -@app.get("/mcp/resource") -async def get_mcp_resource(uri: str): - """Get an MCP resource by URI""" - try: - from .mcp_integration import FuzeAgentMCPServer - - if not uri: - raise HTTPException(status_code=400, detail="uri parameter is required") - - mcp_server = FuzeAgentMCPServer() - resource = await mcp_server.handle_resource_request(uri) - - return resource - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get MCP resource: {str(e)}" - ) - - -@app.get("/tasks/{task_id}/mcp-context") -async def get_task_mcp_context(task_id: str): - """Get MCP context for a task""" - try: - from .mcp_integration import FuzeAgentMCPServer, MCPClaudeIntegration - - execution = app.state.task_execution_engine.active_executions.get(task_id) - if not execution: - raise HTTPException(status_code=404, detail="Task not found or not active") - - mcp_server = FuzeAgentMCPServer() - mcp_integration = MCPClaudeIntegration(mcp_server) - - session_id = execution.claude_session_id or f"session-{task_id}" - context = await mcp_integration.get_session_context( - session_id=session_id, agent_id=execution.agent_id, task_id=task_id - ) - - return context - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get MCP context: {str(e)}" - ) - - -@app.post( - "/agents/{agent_id}/mcp-setup", - tags=["mcp-integration"], - summary="Setup Agent MCP Integration", - description="Configure MCP integration for an AI agent", -) -async def setup_agent_mcp( - agent_id: str = Path(..., description="Agent ID"), - setup_data: AgentMCPSetupRequest = Body(...), -): - """Set up MCP integration for an agent""" - try: - from .mcp_integration import FuzeAgentMCPServer, MCPClaudeIntegration - - task_id = setup_data.get("task_id") - session_id = setup_data.get("session_id") - - if not task_id: - raise HTTPException(status_code=400, detail="task_id is required") - - mcp_server = FuzeAgentMCPServer() - mcp_integration = MCPClaudeIntegration(mcp_server) - - # Set up MCP for Claude session - mcp_config = await mcp_integration.setup_claude_session_mcp( - session_id=session_id or f"session-{task_id}", - agent_id=agent_id, - task_id=task_id, - ) - - return { - "agent_id": agent_id, - "task_id": task_id, - "mcp_config": mcp_config, - "status": "mcp_configured", - } - - except Exception as e: - raise HTTPException(status_code=500, detail=f"Failed to setup MCP: {str(e)}") - - -# Multi-Agent Coordination Endpoints -@app.post( - "/tasks/{task_id}/coordinate", - tags=["multi-agent-coordination"], - summary="Initiate Multi-Agent Coordination", - description="Initiate multi-agent coordination for complex tasks", - response_model=CoordinationResponse, -) -async def initiate_task_coordination( - task_id: str = Path(..., description="Task ID to coordinate"), - coordination_request: CoordinationRequest = Body(...), -): - """Initiate multi-agent coordination for a complex task""" - try: - from .multi_agent_coordinator import CoordinationMode - - coordination_mode = coordination_request.get( - "coordination_mode", "collaborative" - ) - required_agents = coordination_request.get("required_agents") - required_skills = coordination_request.get("required_skills") - - # Validate coordination mode - try: - coord_mode = CoordinationMode(coordination_mode) - except ValueError: - raise HTTPException( - status_code=400, - detail=f"Invalid coordination mode: {coordination_mode}", - ) - - # Get multi-agent coordinator - coordinator = getattr( - app.state.task_execution_engine, "multi_agent_coordinator", None - ) - if not coordinator: - raise HTTPException( - status_code=503, detail="Multi-agent coordination not available" - ) - - # Initiate coordination - session_id = await coordinator.initiate_coordination( - task_id=task_id, - coordination_mode=coord_mode, - required_agents=required_agents, - required_skills=required_skills, - ) - - if session_id: - return { - "task_id": task_id, - "coordination_session_id": session_id, - "status": "coordination_initiated", - "coordination_mode": coordination_mode, - } - else: - return { - "task_id": task_id, - "status": "coordination_not_needed", - "message": "Task does not require multi-agent coordination", - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to initiate coordination: {str(e)}" - ) - - -@app.get("/coordination/{session_id}") -async def get_coordination_status(session_id: str): - """Get status of a coordination session""" - try: - coordinator = getattr( - app.state.task_execution_engine, "multi_agent_coordinator", None - ) - if not coordinator: - raise HTTPException( - status_code=503, detail="Multi-agent coordination not available" - ) - - status = await coordinator.get_coordination_status(session_id) - - if status: - return status - else: - raise HTTPException( - status_code=404, detail="Coordination session not found" - ) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get coordination status: {str(e)}" - ) - - -@app.post("/coordination/{session_id}/cancel") -async def cancel_coordination(session_id: str): - """Cancel a coordination session""" - try: - coordinator = getattr( - app.state.task_execution_engine, "multi_agent_coordinator", None - ) - if not coordinator: - raise HTTPException( - status_code=503, detail="Multi-agent coordination not available" - ) - - success = await coordinator.cancel_coordination(session_id) - - if success: - return {"coordination_session_id": session_id, "status": "cancelled"} - else: - raise HTTPException( - status_code=404, detail="Coordination session not found" - ) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to cancel coordination: {str(e)}" - ) - - -@app.post("/agents/{from_agent_id}/communicate/{to_agent_id}") -async def send_agent_communication( - from_agent_id: str, to_agent_id: str, communication_data: dict -): - """Send communication between agents""" - try: - message_type = communication_data.get("message_type", "notification") - content = communication_data.get("content", "") - metadata = communication_data.get("metadata", {}) - - if not content: - raise HTTPException(status_code=400, detail="Content cannot be empty") - - coordinator = getattr( - app.state.task_execution_engine, "multi_agent_coordinator", None - ) - if not coordinator: - raise HTTPException( - status_code=503, detail="Multi-agent coordination not available" - ) - - communication_id = await coordinator.send_agent_communication( - from_agent_id=from_agent_id, - to_agent_id=to_agent_id, - message_type=message_type, - content=content, - metadata=metadata, - ) - - return { - "communication_id": communication_id, - "from_agent_id": from_agent_id, - "to_agent_id": to_agent_id, - "status": "sent", - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to send agent communication: {str(e)}" - ) - - -@app.get("/coordination/active") -async def get_active_coordinations(): - """Get all active coordination sessions""" - try: - coordinator = getattr( - app.state.task_execution_engine, "multi_agent_coordinator", None - ) - if not coordinator: - raise HTTPException( - status_code=503, detail="Multi-agent coordination not available" - ) - - active_sessions = [] - for session_id in coordinator.active_sessions.keys(): - status = await coordinator.get_coordination_status(session_id) - if status: - active_sessions.append(status) - - return {"active_coordinations": active_sessions, "count": len(active_sessions)} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get active coordinations: {str(e)}" - ) - - -# WebSocket for coordination updates -@app.websocket("/ws/coordination/{session_id}") -async def coordination_websocket_endpoint(websocket: WebSocket, session_id: str): - """WebSocket endpoint for real-time coordination updates""" - await websocket.accept() - try: - coordinator = getattr( - app.state.task_execution_engine, "multi_agent_coordinator", None - ) - if not coordinator: - await websocket.send_json( - {"type": "error", "message": "Multi-agent coordination not available"} - ) - await websocket.close() - return - - # Monitor coordination session - while True: - try: - status = await coordinator.get_coordination_status(session_id) - if status: - await websocket.send_json( - { - "type": "coordination_update", - "session_id": session_id, - "data": status, - "timestamp": datetime.now().isoformat(), - } - ) - - # If coordination is completed or failed, send final update - if status.get("status") in ["completed", "failed", "cancelled"]: - await websocket.send_json( - { - "type": "coordination_finished", - "session_id": session_id, - "final_status": status.get("status"), - "timestamp": datetime.now().isoformat(), - } - ) - break - else: - await websocket.send_json( - { - "type": "error", - "message": f"Coordination session {session_id} not found", - } - ) - break - - await asyncio.sleep(3) # Update every 3 seconds - - except Exception as e: - await websocket.send_json( - { - "type": "error", - "message": f"Error monitoring coordination: {str(e)}", - } - ) - - except Exception as e: - print(f"Coordination WebSocket error for {session_id}: {e}") - finally: - await websocket.close() - - -# --------------------------------------------------------------------------- -# Agent relay WebSocket (Track 4) -# --------------------------------------------------------------------------- -@app.websocket("/agent-relay/{agent_id}") -async def agent_relay_endpoint(websocket: WebSocket, agent_id: str): - """ - Agent pods connect here to stream their session output. - Dashboard clients connect here to watch a specific agent's session. - Both use the same endpoint — first JSON message determines role: - {"role": "agent"} -> agent pod streaming output - {"role": "subscriber"} -> human dashboard watcher (default) - """ - await websocket.accept() - role = None - try: - init_msg = await websocket.receive_json() - role = init_msg.get("role", "subscriber") - - if role == "agent": - # Stream from agent pod to all subscribers - async for data in websocket.iter_json(): - msg = {"agentId": agent_id, **data} - dead = [] - for sub in list(agent_relay_subscribers[agent_id]): - try: - await sub.send_json(msg) - except Exception: - dead.append(sub) - for d in dead: - agent_relay_subscribers[agent_id].remove(d) - else: - # Human dashboard subscriber — wait for messages from agent - agent_relay_subscribers[agent_id].append(websocket) - await websocket.receive_text() # keep alive until disconnect - except WebSocketDisconnect: - pass - except Exception as e: - logger.warning(f"agent-relay {agent_id}: {e}") - finally: - subs = agent_relay_subscribers.get(agent_id, []) - if role != "agent" and websocket in subs: - subs.remove(websocket) - - -# Model Configuration and API Key Management Endpoints -@app.post( - "/organizations/{organization_id}/providers/{provider}/credentials", - tags=["model-configuration"], - summary="Store Provider API Credentials", - description="Store encrypted API credentials for a model provider", -) -async def store_provider_credentials( - organization_id: str = Path(..., description="Organization ID"), - provider: str = Path(..., description="Provider name"), - credentials: ProviderCredentialsRequest = Body(...), -): - """Store encrypted API credentials for a model provider at organization level""" - try: - from .model_configuration import ModelProvider, model_config_manager - - # Validate provider - try: - provider_enum = ModelProvider(provider) - except ValueError: - raise HTTPException( - status_code=400, detail=f"Unsupported provider: {provider}" - ) - - success = await model_config_manager.store_provider_credentials( - organization_id=organization_id, - provider=provider_enum, - api_key=credentials.api_key, - endpoint_url=credentials.endpoint_url, - additional_config=credentials.additional_config, - ) - - if success: - return { - "organization_id": organization_id, - "provider": provider, - "status": "credentials_stored", - "message": "API credentials stored successfully", - } - else: - raise HTTPException(status_code=500, detail="Failed to store credentials") - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to store provider credentials: {str(e)}" - ) - - -@app.get( - "/organizations/{organization_id}/models", - tags=["model-configuration"], - summary="Get Available Models", - description="Get available AI models for an organization", -) -async def get_available_models( - organization_id: str = Path(..., description="Organization ID"), - provider: Optional[str] = Query(None, description="Filter by provider"), - capabilities: Optional[str] = Query( - None, description="Filter by capabilities (comma-separated)" - ), -): - """Get available AI models with provider credential validation""" - try: - from .model_configuration import ( - ModelCapability, - ModelProvider, - model_config_manager, - ) - - provider_filter = None - if provider: - try: - provider_filter = ModelProvider(provider) - except ValueError: - raise HTTPException( - status_code=400, detail=f"Invalid provider: {provider}" - ) - - capabilities_filter = None - if capabilities: - try: - capabilities_filter = [ - ModelCapability(cap.strip()) for cap in capabilities.split(",") - ] - except ValueError as e: - raise HTTPException( - status_code=400, detail=f"Invalid capability: {str(e)}" - ) - - models = await model_config_manager.get_available_models( - organization_id=organization_id, - provider=provider_filter, - capabilities=capabilities_filter, - ) - - return { - "organization_id": organization_id, - "models": models, - "count": len(models), - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get available models: {str(e)}" - ) - - -@app.post( - "/agents/{agent_id}/model-configuration", - tags=["model-configuration"], - summary="Configure Agent Model Settings", - description="Configure model settings and preferences for an agent", -) -async def configure_agent_model( - agent_id: str = Path(..., description="Agent ID"), - config: AgentModelConfigRequest = Body(...), -): - """Configure model settings for an AI agent""" - try: - from .model_configuration import AgentModelConfig, model_config_manager - - agent_config = AgentModelConfig( - agent_id=agent_id, - primary_model=config.primary_model, - fallback_models=config.fallback_models, - temperature=config.temperature, - max_tokens=config.max_tokens, - top_p=config.top_p, - frequency_penalty=config.frequency_penalty, - presence_penalty=config.presence_penalty, - custom_instructions=config.custom_instructions, - use_function_calling=config.use_function_calling, - streaming_enabled=config.streaming_enabled, - cost_limit_per_task=config.cost_limit_per_task, - ) - - success = await model_config_manager.configure_agent_model( - agent_id, agent_config - ) - - if success: - return { - "agent_id": agent_id, - "status": "configured", - "primary_model": config.primary_model, - "fallback_models": config.fallback_models, - } - else: - raise HTTPException( - status_code=500, detail="Failed to configure agent model" - ) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to configure agent model: {str(e)}" - ) - - -@app.get( - "/agents/{agent_id}/model-configuration", - tags=["model-configuration"], - summary="Get Agent Model Configuration", - description="Get current model configuration for an agent", -) -async def get_agent_model_configuration( - agent_id: str = Path(..., description="Agent ID") -): - """Get model configuration for an AI agent""" - try: - from .model_configuration import model_config_manager - - config = await model_config_manager.get_agent_model_config(agent_id) - - if config: - return { - "agent_id": agent_id, - "configuration": { - "primary_model": config.primary_model, - "fallback_models": config.fallback_models, - "temperature": config.temperature, - "max_tokens": config.max_tokens, - "top_p": config.top_p, - "frequency_penalty": config.frequency_penalty, - "presence_penalty": config.presence_penalty, - "custom_instructions": config.custom_instructions, - "use_function_calling": config.use_function_calling, - "streaming_enabled": config.streaming_enabled, - "cost_limit_per_task": config.cost_limit_per_task, - "created_at": config.created_at.isoformat(), - "updated_at": config.updated_at.isoformat(), - }, - } - else: - raise HTTPException( - status_code=404, detail="Agent model configuration not found" - ) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get agent model configuration: {str(e)}" - ) - - -@app.post( - "/agents/{agent_id}/tasks/cost-estimate", - tags=["model-configuration"], - summary="Estimate Task Cost", - description="Estimate the cost of executing a task with the agent's model configuration", -) -async def estimate_task_cost( - agent_id: str = Path(..., description="Agent ID"), - request: TaskCostEstimateRequest = Body(...), -): - """Estimate cost for task execution based on agent's model configuration""" - try: - from .model_configuration import model_config_manager - - estimate = await model_config_manager.estimate_task_cost( - agent_id=agent_id, - task_description=request.task_description, - estimated_complexity=request.estimated_complexity, - ) - - return estimate - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to estimate task cost: {str(e)}" - ) - - -@app.get( - "/organizations/{organization_id}/model-usage", - tags=["model-configuration"], - summary="Get Model Usage Statistics", - description="Get model usage statistics and costs for an organization", -) -async def get_organization_model_usage( - organization_id: str = Path(..., description="Organization ID"), - days: int = Query(30, ge=1, le=365, description="Number of days to analyze"), -): - """Get model usage statistics and costs for an organization""" - try: - from .model_configuration import model_config_manager - - usage = await model_config_manager.get_organization_model_usage( - organization_id=organization_id, days=days - ) - - return usage - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get model usage: {str(e)}" - ) - - -@app.get( - "/agents/{agent_id}/model-recommendations", - tags=["model-configuration"], - summary="Get Model Recommendations", - description="Get model recommendations for an agent based on task capabilities", -) -async def get_model_recommendations( - agent_id: str = Path(..., description="Agent ID"), - capabilities: str = Query( - ..., description="Required capabilities (comma-separated)" - ), - cost_limit: Optional[float] = Query( - None, ge=0.0, description="Maximum cost limit in USD" - ), -): - """Get model recommendations based on task capabilities and cost constraints""" - try: - from .model_configuration import ModelCapability, model_config_manager - - # Parse capabilities - try: - capability_list = [ - ModelCapability(cap.strip()) for cap in capabilities.split(",") - ] - except ValueError as e: - raise HTTPException(status_code=400, detail=f"Invalid capability: {str(e)}") - - recommended_model = await model_config_manager.get_model_for_task( - agent_id=agent_id, task_capabilities=capability_list, cost_limit=cost_limit - ) - - if recommended_model: - return { - "agent_id": agent_id, - "recommended_model": recommended_model, - "capabilities": capabilities, - "cost_limit": cost_limit, - } - else: - return { - "agent_id": agent_id, - "recommended_model": None, - "message": "No suitable model found for the specified requirements", - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get model recommendations: {str(e)}" - ) - - -# Knowledge Management and Notification Endpoints - - -@app.get( - "/knowledge/notifications/{recipient_type}/{recipient_id}", - tags=["knowledge-management"], - summary="Get Knowledge Notifications", - description="Get notifications about knowledge updates, conflicts, and opportunities", -) -async def get_knowledge_notifications( - recipient_type: str = Path( - ..., description="Recipient type (agent, team, organization)" - ), - recipient_id: str = Path(..., description="Recipient ID"), - limit: int = Query(20, ge=1, le=100, description="Maximum notifications to return"), - status_filter: Optional[str] = Query( - None, description="Filter by status (unread, read, acknowledged)" - ), - notification_type_filter: Optional[str] = Query( - None, description="Filter by type (comma-separated)" - ), -): - """Get knowledge notifications for a recipient""" - try: - from .knowledge_notification_service import ( - KnowledgeNotificationService, - NotificationStatus, - NotificationType, - ) - - # Initialize notification service if not already done - if not hasattr(app.state, "notification_service"): - database_url = os.getenv( - "DATABASE_URL", - "postgresql://postgres:password@postgres:5432/ai_context", - ) - app.state.notification_service = KnowledgeNotificationService(database_url) - await app.state.notification_service.initialize() - - # Parse filters - status_filters = None - if status_filter: - try: - status_filters = [ - NotificationStatus(s.strip()) for s in status_filter.split(",") - ] - except ValueError as e: - raise HTTPException( - status_code=400, detail=f"Invalid status filter: {str(e)}" - ) - - type_filters = None - if notification_type_filter: - try: - type_filters = [ - NotificationType(t.strip()) - for t in notification_type_filter.split(",") - ] - except ValueError as e: - raise HTTPException( - status_code=400, - detail=f"Invalid notification type filter: {str(e)}", - ) - - notifications = ( - await app.state.notification_service.get_notifications_for_recipient( - recipient_type=recipient_type, - recipient_id=recipient_id, - limit=limit, - status_filter=status_filters, - notification_type_filter=type_filters, - ) - ) - - return { - "recipient_type": recipient_type, - "recipient_id": recipient_id, - "notifications": [ - { - "id": n.id, - "notification_type": n.notification_type.value, - "title": n.title, - "message": n.message, - "knowledge_id": n.knowledge_id, - "knowledge_type": n.knowledge_type, - "priority": n.priority.value, - "requires_action": n.requires_action, - "status": n.status.value, - "suggested_actions": n.suggested_actions, - "metadata": n.metadata, - "created_at": n.created_at.isoformat(), - "expires_at": n.expires_at.isoformat() if n.expires_at else None, - } - for n in notifications - ], - "count": len(notifications), - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get knowledge notifications: {str(e)}" - ) - - -@app.put( - "/knowledge/notifications/{notification_id}/status", - tags=["knowledge-management"], - summary="Update Notification Status", - description="Mark notification as read, acknowledged, or acted upon", -) -async def update_notification_status( - notification_id: str = Path(..., description="Notification ID"), - status: str = Body(..., description="New notification status"), - action_taken: Optional[Dict[str, Any]] = Body( - None, description="Optional action taken metadata" - ), -): - """Update notification status and optional action taken""" - try: - from .knowledge_notification_service import NotificationStatus - - # Validate status - try: - notification_status = NotificationStatus(status) - except ValueError: - raise HTTPException(status_code=400, detail=f"Invalid status: {status}") - - success = await app.state.notification_service.mark_notification_status( - notification_id=notification_id, - status=notification_status, - action_taken=action_taken, - ) - - if success: - return { - "notification_id": notification_id, - "status": status, - "updated": True, - } - else: - raise HTTPException(status_code=404, detail="Notification not found") - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to update notification status: {str(e)}" - ) - - -@app.get( - "/knowledge/notifications/statistics", - tags=["knowledge-management"], - summary="Get Notification Statistics", - description="Get comprehensive notification statistics and analytics", -) -async def get_notification_statistics( - organization_id: Optional[str] = Query( - None, description="Filter by organization ID" - ), - days_back: int = Query(30, ge=1, le=365, description="Days of history to analyze"), -): - """Get notification statistics and analytics""" - try: - stats = await app.state.notification_service.get_notification_statistics( - organization_id=organization_id, days_back=days_back - ) - - return stats - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get notification statistics: {str(e)}" - ) - - -@app.post( - "/knowledge/organizations/{organization_id}/add", - tags=["knowledge-management"], - summary="Add Organizational Knowledge", - description="Add knowledge to organization-level knowledge base", -) -async def add_organizational_knowledge( - organization_id: str = Path(..., description="Organization ID"), - title: str = Body(..., description="Knowledge title"), - content: str = Body(..., description="Knowledge content"), - content_type: str = Body("documentation", description="Content type"), - knowledge_category: str = Body("development", description="Knowledge category"), - source_agent_id: Optional[str] = Body(None, description="Source agent ID"), - source_team_id: Optional[str] = Body(None, description="Source team ID"), - tags: List[str] = Body(default_factory=list, description="Knowledge tags"), - metadata: Dict[str, Any] = Body( - default_factory=dict, description="Additional metadata" - ), -): - """Add knowledge to organization-level knowledge base""" - try: - from .organization_rag_manager import ( - ContentType, - KnowledgeCategory, - OrganizationRAGManager, - SourceType, - ) - - # Initialize services if not already done - if not hasattr(app.state, "org_rag_manager"): - database_url = os.getenv( - "DATABASE_URL", - "postgresql://postgres:password@postgres:5432/ai_context", - ) - app.state.org_rag_manager = OrganizationRAGManager(database_url) - await app.state.org_rag_manager.initialize() - - # Validate enums - try: - content_type_enum = ContentType(content_type) - category_enum = KnowledgeCategory(knowledge_category) - except ValueError as e: - raise HTTPException(status_code=400, detail=f"Invalid enum value: {str(e)}") - - knowledge_id = await app.state.org_rag_manager.add_knowledge( - organization_id=organization_id, - title=title, - content=content, - content_type=content_type_enum, - knowledge_category=category_enum, - source_type=SourceType.MANUAL_INPUT, - source_agent_id=source_agent_id, - source_team_id=source_team_id, - tags=tags, - metadata=metadata, - ) - - return { - "knowledge_id": knowledge_id, - "organization_id": organization_id, - "title": title, - "status": "added", - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to add organizational knowledge: {str(e)}" - ) - - -@app.get( - "/knowledge/organizations/{organization_id}/search", - tags=["knowledge-management"], - summary="Search Organizational Knowledge", - description="Search organization-level knowledge base", -) -async def search_organizational_knowledge( - organization_id: str = Path(..., description="Organization ID"), - query: str = Query(..., description="Search query"), - limit: int = Query(10, ge=1, le=50, description="Maximum results to return"), - min_similarity: float = Query( - 0.3, ge=0.0, le=1.0, description="Minimum similarity threshold" - ), - categories: Optional[str] = Query( - None, description="Filter by categories (comma-separated)" - ), -): - """Search organization-level knowledge base""" - try: - from .organization_rag_manager import KnowledgeCategory - - # Parse categories - category_filters = None - if categories: - try: - category_filters = [ - KnowledgeCategory(cat.strip()) for cat in categories.split(",") - ] - except ValueError as e: - raise HTTPException( - status_code=400, detail=f"Invalid category: {str(e)}" - ) - - search_results = await app.state.org_rag_manager.search_knowledge( - organization_id=organization_id, - query=query, - limit=limit, - min_similarity=min_similarity, - categories=category_filters, - ) - - results = [] - for result in search_results: - results.append( - { - "knowledge_id": result.knowledge.id, - "title": result.knowledge.title, - "content_preview": ( - result.knowledge.content[:200] + "..." - if len(result.knowledge.content) > 200 - else result.knowledge.content - ), - "category": result.knowledge.knowledge_category.value, - "content_type": result.knowledge.content_type.value, - "similarity_score": result.similarity_score, - "combined_score": result.combined_score, - "quality_score": result.knowledge.quality_score, - "usage_count": result.knowledge.usage_count, - "created_at": result.knowledge.created_at.isoformat(), - "tags": result.knowledge.tags, - "metadata": result.knowledge.metadata, - } - ) - - return { - "organization_id": organization_id, - "query": query, - "results": results, - "count": len(results), - } - - except Exception as e: - raise HTTPException( - status_code=500, - detail=f"Failed to search organizational knowledge: {str(e)}", - ) - - -@app.get( - "/knowledge/context-enhancement/{agent_id}", - tags=["knowledge-management"], - summary="Get Enhanced Context for Agent", - description="Get enhanced context with relevant organizational knowledge for task execution", -) -async def get_enhanced_context_for_agent( - agent_id: str = Path(..., description="Agent ID"), - task_description: str = Query( - ..., description="Task description for context enhancement" - ), - task_type: Optional[str] = Query(None, description="Task type"), - technologies: Optional[str] = Query( - None, description="Technologies involved (comma-separated)" - ), -): - """Get enhanced context with relevant knowledge for agent task execution""" - try: - from .context_enhancement_service import ContextEnhancementService - - # Initialize context enhancement service if needed - if not hasattr(app.state, "context_enhancement_service"): - database_url = os.getenv( - "DATABASE_URL", - "postgresql://postgres:password@postgres:5432/ai_context", - ) - # These would be initialized in the lifespan - if hasattr(app.state, "org_rag_manager") and hasattr( - app.state, "team_knowledge_manager" - ): - app.state.context_enhancement_service = ContextEnhancementService( - database_url=database_url, - org_rag_manager=app.state.org_rag_manager, - team_knowledge_manager=app.state.team_knowledge_manager, - ) - await app.state.context_enhancement_service.initialize() - else: - raise HTTPException( - status_code=503, - detail="Knowledge management services not initialized", - ) - - # Build task data - task_data = { - "description": task_description, - "task_type": task_type, - "technologies": technologies.split(",") if technologies else [], - } - - enhanced_context = ( - await app.state.context_enhancement_service.enhance_agent_context( - agent_id=agent_id, task_data=task_data - ) - ) - - return { - "agent_id": agent_id, - "task_description": task_description, - "enhanced_context": { - "organizational_knowledge_count": len( - enhanced_context.organizational_knowledge - ), - "team_knowledge_count": len(enhanced_context.team_knowledge), - "similar_task_insights_count": len( - enhanced_context.similar_task_insights - ), - "success_patterns": enhanced_context.success_patterns, - "common_pitfalls": enhanced_context.common_pitfalls, - "recommended_approaches": enhanced_context.recommended_approaches, - "context_summary": enhanced_context.context_summary, - "enhancement_metadata": enhanced_context.enhancement_metadata, - }, - "organizational_knowledge": [ - { - "knowledge_id": item.knowledge_id, - "title": item.title, - "category": item.category, - "relevance_score": item.relevance_score, - "confidence_score": item.confidence_score, - "content_preview": ( - item.content[:200] + "..." - if len(item.content) > 200 - else item.content - ), - } - for item in enhanced_context.organizational_knowledge - ], - "team_knowledge": [ - { - "knowledge_id": item.knowledge_id, - "title": item.title, - "category": item.category, - "relevance_score": item.relevance_score, - "confidence_score": item.confidence_score, - "content_preview": ( - item.content[:200] + "..." - if len(item.content) > 200 - else item.content - ), - } - for item in enhanced_context.team_knowledge - ], - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get enhanced context: {str(e)}" - ) - - -@app.get( - "/knowledge/analytics/organizations/{organization_id}/insights", - tags=["knowledge-management"], - summary="Get Organizational Knowledge Insights", - description="Get comprehensive analytics and insights about organizational knowledge", -) -async def get_organizational_knowledge_insights( - organization_id: str = Path(..., description="Organization ID"), - analysis_period_days: int = Query( - 30, ge=7, le=365, description="Analysis period in days" - ), -): - """Get comprehensive organizational knowledge insights and analytics""" - try: - insights = ( - await app.state.knowledge_analytics_service.get_organizational_insights( - organization_id=organization_id, - analysis_period_days=analysis_period_days, - ) - ) - - return { - "organization_id": organization_id, - "analysis_period_days": analysis_period_days, - "insights": { - "total_knowledge_items": insights.total_knowledge_items, - "knowledge_growth_rate": insights.knowledge_growth_rate, - "knowledge_utilization_rate": insights.knowledge_utilization_rate, - "knowledge_freshness_score": insights.knowledge_freshness_score, - "cross_team_sharing_rate": insights.cross_team_sharing_rate, - "propagation_efficiency": insights.propagation_efficiency, - "top_performing_categories": insights.top_performing_categories, - "knowledge_gaps": insights.knowledge_gaps, - "agent_knowledge_engagement": insights.agent_knowledge_engagement, - "team_knowledge_contribution": insights.team_knowledge_contribution, - "recommendations": insights.recommendations, - }, - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get organizational insights: {str(e)}" - ) - - -@app.get( - "/knowledge/analytics/organizations/{organization_id}/effectiveness", - tags=["knowledge-management"], - summary="Analyze Knowledge Effectiveness", - description="Analyze effectiveness and performance of knowledge items", -) -async def analyze_knowledge_effectiveness( - organization_id: str = Path(..., description="Organization ID"), - knowledge_category: Optional[str] = Query( - None, description="Filter by knowledge category" - ), - min_usage_count: int = Query( - 3, ge=1, description="Minimum usage count for analysis" - ), -): - """Analyze effectiveness of knowledge items in the organization""" - try: - effectiveness_metrics = ( - await app.state.knowledge_analytics_service.analyze_knowledge_effectiveness( - organization_id=organization_id, - knowledge_category=knowledge_category, - min_usage_count=min_usage_count, - ) - ) - - results = [] - for metric in effectiveness_metrics: - results.append( - { - "knowledge_id": metric.knowledge_id, - "title": metric.title, - "category": metric.category, - "usage_count": metric.usage_count, - "success_correlation": metric.success_correlation, - "average_relevance": metric.average_relevance, - "agent_adoption_rate": metric.agent_adoption_rate, - "team_adoption_rate": metric.team_adoption_rate, - "quality_score": metric.quality_score, - "recency_score": metric.recency_score, - "overall_effectiveness": metric.overall_effectiveness, - "trend_direction": metric.trend_direction, - "optimization_suggestions": metric.optimization_suggestions, - } - ) - - return { - "organization_id": organization_id, - "effectiveness_analysis": results, - "total_analyzed": len(results), - "summary": { - "avg_effectiveness": sum(r["overall_effectiveness"] for r in results) - / max(len(results), 1), - "top_performers": sorted( - results, key=lambda x: x["overall_effectiveness"], reverse=True - )[:5], - "needs_attention": [ - r for r in results if r["overall_effectiveness"] < 0.5 - ], - }, - } - - except Exception as e: - raise HTTPException( - status_code=500, - detail=f"Failed to analyze knowledge effectiveness: {str(e)}", - ) - - -@app.get( - "/knowledge/analytics/agents/{agent_id}/profile", - tags=["knowledge-management"], - summary="Get Agent Knowledge Profile", - description="Get detailed knowledge profile and analytics for an agent", -) -async def get_agent_knowledge_profile( - agent_id: str = Path(..., description="Agent ID"), - analysis_period_days: int = Query( - 60, ge=7, le=365, description="Analysis period in days" - ), -): - """Get detailed knowledge profile for an agent""" - try: - profile = ( - await app.state.knowledge_analytics_service.get_agent_knowledge_profile( - agent_id=agent_id, analysis_period_days=analysis_period_days - ) - ) - - if not profile: - raise HTTPException( - status_code=404, detail="Agent not found or no knowledge data available" - ) - - return { - "agent_id": agent_id, - "analysis_period_days": analysis_period_days, - "profile": { - "agent_name": profile.agent_name, - "team_id": profile.team_id, - "knowledge_consumption_rate": profile.knowledge_consumption_rate, - "knowledge_creation_rate": profile.knowledge_creation_rate, - "expertise_areas": profile.expertise_areas, - "knowledge_application_success": profile.knowledge_application_success, - "learning_velocity": profile.learning_velocity, - "knowledge_sharing_activity": profile.knowledge_sharing_activity, - "preferred_knowledge_types": profile.preferred_knowledge_types, - "knowledge_gaps": profile.knowledge_gaps, - "optimization_recommendations": profile.optimization_recommendations, - }, - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get agent knowledge profile: {str(e)}" - ) - - -@app.get( - "/knowledge/analytics/organizations/{organization_id}/optimization", - tags=["knowledge-management"], - summary="Get Knowledge Optimization Recommendations", - description="Get comprehensive recommendations for knowledge system optimization", -) -async def get_knowledge_optimization_recommendations( - organization_id: str = Path(..., description="Organization ID"), - focus_area: Optional[str] = Query( - None, - description="Focus area (utilization, quality, gaps, propagation, collaboration)", - ), -): - """Generate comprehensive knowledge optimization recommendations""" - try: - recommendations = await app.state.knowledge_analytics_service.generate_knowledge_optimization_recommendations( - organization_id=organization_id, focus_area=focus_area - ) - - return { - "organization_id": organization_id, - "focus_area": focus_area, - "recommendations": recommendations, - "total_recommendations": len(recommendations), - } - - except Exception as e: - raise HTTPException( - status_code=500, - detail=f"Failed to get optimization recommendations: {str(e)}", - ) - - -@app.get( - "/knowledge/analytics/organizations/{organization_id}/trends", - tags=["knowledge-management"], - summary="Get Knowledge Trends Analysis", - description="Analyze knowledge trends and patterns over time", -) -async def get_knowledge_trends_analysis( - organization_id: str = Path(..., description="Organization ID"), - trend_period_days: int = Query( - 90, ge=30, le=365, description="Trend analysis period in days" - ), -): - """Get comprehensive knowledge trends analysis""" - try: - trends = ( - await app.state.knowledge_analytics_service.get_knowledge_trends_analysis( - organization_id=organization_id, trend_period_days=trend_period_days - ) - ) - - return { - "organization_id": organization_id, - "trend_period_days": trend_period_days, - "trends": trends, - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get knowledge trends: {str(e)}" - ) - - -# Memory-Enhanced Agents Endpoints - - -@app.post( - "/agents/{agent_id}/deploy-memory", - tags=["memory-agents"], - summary="Deploy Memory-Enabled Agent", - description="Deploy an agent with persistent memory capabilities", -) -async def deploy_memory_enabled_agent( - agent_id: str = Path(..., description="Agent ID"), - template_id: str = Body(..., description="Agent template ID"), - task_id: Optional[str] = Body(None, description="Optional specific task ID"), - repository_settings: Optional[Dict[str, Any]] = Body( - None, description="Repository settings" - ), -): - """Deploy a memory-enabled autonomous agent container""" - try: - result = await app.state.agent_manager.deploy_memory_enabled_agent( - agent_id=agent_id, - template_id=template_id, - task_id=task_id, - repository_settings=repository_settings, - ) - - if result["success"]: - return result - else: - raise HTTPException(status_code=500, detail=result["error"]) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to deploy memory-enabled agent: {str(e)}" - ) - - -@app.get( - "/agents/{agent_id}/memory-status", - tags=["memory-agents"], - summary="Get Agent Memory Status", - description="Get agent memory status and expertise summary", -) -async def get_agent_memory_status(agent_id: str = Path(..., description="Agent ID")): - """Get agent memory status, expertise metrics, and insights""" - try: - status = await app.state.agent_manager.get_agent_memory_status(agent_id) - return status - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get agent memory status: {str(e)}" - ) - - -@app.post( - "/agents/{agent_id}/memory-tasks", - tags=["memory-agents"], - summary="Assign Task to Memory Agent", - description="Assign a task to a memory-enabled agent", -) -async def assign_task_to_memory_agent( - agent_id: str = Path(..., description="Agent ID"), - task_id: str = Body(..., description="Task ID"), - task_data: Dict[str, Any] = Body(..., description="Task data"), -): - """Assign a task to a memory-enabled agent for autonomous execution""" - try: - result = await app.state.agent_manager.assign_task_to_memory_agent( - agent_id=agent_id, task_id=task_id, task_data=task_data - ) - - if result["success"]: - return result - else: - raise HTTPException(status_code=400, detail=result["error"]) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to assign task to memory agent: {str(e)}" - ) - - -@app.delete( - "/agents/{agent_id}/memory", - tags=["memory-agents"], - summary="Stop Memory-Enabled Agent", - description="Stop a memory-enabled agent container", -) -async def stop_memory_enabled_agent(agent_id: str = Path(..., description="Agent ID")): - """Stop and clean up a memory-enabled agent container""" - try: - result = await app.state.agent_manager.stop_memory_enabled_agent(agent_id) - - if result["success"]: - return result - else: - raise HTTPException(status_code=400, detail=result["error"]) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to stop memory-enabled agent: {str(e)}" - ) - - -@app.get( - "/system/expertise-dashboard", - tags=["memory-agents"], - summary="Get System Expertise Dashboard", - description="Get system-wide expertise and memory analytics", -) -async def get_system_expertise_dashboard(): - """Get comprehensive dashboard of system expertise and memory analytics""" - try: - dashboard = await app.state.agent_manager.get_system_expertise_dashboard() - return dashboard - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get expertise dashboard: {str(e)}" - ) - - -@app.get( - "/agents/{agent_id}/tasks/pending", - tags=["memory-agents"], - summary="Get Pending Tasks for Agent", - description="Get pending tasks for a memory-enabled agent", -) -async def get_pending_tasks_for_agent( - agent_id: str = Path(..., description="Agent ID"), - limit: int = Query( - 10, ge=1, le=50, description="Maximum number of tasks to return" - ), -): - """Get pending tasks that a memory-enabled agent can pick up""" - try: - async with get_db_connection() as conn: - tasks = await conn.fetch( - """ - SELECT id, title, description, type, complexity, language, - framework, requirements, created_at - FROM tasks - WHERE agent_id = $1 - AND status = 'pending' - AND assigned_to_memory_agent = true - ORDER BY created_at ASC - LIMIT $2 - """, - agent_id, - limit, - ) - - return { - "agent_id": agent_id, - "tasks": [dict(task) for task in tasks], - "count": len(tasks), - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get pending tasks: {str(e)}" - ) - - -@app.put( - "/tasks/{task_id}/status", - tags=["memory-agents"], - summary="Update Task Status", - description="Update task status (used by memory-enabled agents)", -) -async def update_task_status( - task_id: str = Path(..., description="Task ID"), - status: str = Body(..., description="New task status"), - result: Optional[Dict[str, Any]] = Body(None, description="Task result data"), - updated_by: Optional[str] = Body(None, description="ID of agent updating the task"), - container_instance_id: Optional[str] = Body( - None, description="Container instance ID" - ), - updated_at: Optional[str] = Body(None, description="Update timestamp"), -): - """Update task status - used by memory-enabled agents to report progress""" - try: - async with get_db_connection() as conn: - await conn.execute( - """ - UPDATE tasks - SET status = $2, - result = COALESCE($3, result), - updated_by = COALESCE($4, updated_by), - updated_at = NOW() - WHERE id = $1 - """, - task_id, - status, - result, - updated_by, - ) - - # If task is completed, log it for expertise tracking - if status in ["completed", "failed"]: - # The agent's memory system will handle learning from the outcome - pass - - return {"task_id": task_id, "status": status, "updated": True} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to update task status: {str(e)}" - ) - - -@app.post( - "/agents/{agent_id}/register", - tags=["memory-agents"], - summary="Agent Registration", - description="Register agent capabilities and status with orchestrator", -) -async def register_agent_capabilities( - agent_id: str = Path(..., description="Agent ID"), - capabilities: Dict[str, Any] = Body( - ..., description="Agent capabilities and status" - ), -): - """Register or update agent capabilities - used by memory-enabled agents on startup""" - try: - # Update agent capabilities in database - async with get_db_connection() as conn: - await conn.execute( - """ - UPDATE agents - SET config = config || $2, - status = 'active', - updated_at = NOW() - WHERE id = $1 - """, - agent_id, - { - "capabilities": capabilities, - "last_registration": datetime.now().isoformat(), - }, - ) - - # Update in-memory tracking - if agent_id in app.state.agent_manager.memory_enabled_agents: - app.state.agent_manager.memory_enabled_agents[agent_id]["status"] = "active" - - return { - "agent_id": agent_id, - "agent_recognized": True, - "capabilities_accepted": True, - "status": "registered", - } - - except Exception as e: - return { - "agent_id": agent_id, - "agent_recognized": False, - "capabilities_accepted": False, - "error": str(e), - } - - -@app.post( - "/agents/{agent_id}/statistics", - tags=["memory-agents"], - summary="Agent Statistics Update", - description="Update agent performance and memory statistics", -) -async def update_agent_statistics( - agent_id: str = Path(..., description="Agent ID"), - stats: Dict[str, Any] = Body(..., description="Agent statistics"), -): - """Update agent statistics - used by memory-enabled agents for performance tracking""" - try: - # Store statistics for analytics - async with get_db_connection() as conn: - await conn.execute( - """ - UPDATE agents - SET config = config || $2, - updated_at = NOW() - WHERE id = $1 - """, - agent_id, - { - "latest_statistics": stats, - "statistics_updated_at": datetime.now().isoformat(), - }, - ) - - # Clear expertise cache to force refresh - await app.state.agent_manager.expertise_tracker.clear_cache(agent_id) - - return {"agent_id": agent_id, "statistics_updated": True} - - except Exception as e: - return {"agent_id": agent_id, "statistics_updated": False, "error": str(e)} - - -@app.post( - "/agents/{agent_id}/error", - tags=["memory-agents"], - summary="Agent Error Reporting", - description="Report agent errors for monitoring", -) -async def report_agent_error( - agent_id: str = Path(..., description="Agent ID"), - error_data: Dict[str, Any] = Body(..., description="Error information"), -): - """Report agent errors - used by memory-enabled agents for error tracking""" - try: - # Log error for monitoring - logger.error(f"Agent {agent_id} reported error: {error_data}") - - # Update agent status if it's a critical error - if error_data.get("critical", False): - async with get_db_connection() as conn: - await conn.execute( - """ - UPDATE agents - SET status = 'error', - config = config || $2, - updated_at = NOW() - WHERE id = $1 - """, - agent_id, - { - "last_error": error_data, - "error_reported_at": datetime.now().isoformat(), - }, - ) - - return {"agent_id": agent_id, "error_logged": True} - - except Exception as e: - logger.error(f"Failed to log agent error: {e}") - return {"agent_id": agent_id, "error_logged": False} - - -# ============================================================================ -# Goals Management API Endpoints -# ============================================================================ - - -@app.post( - "/organizations/{organization_id}/goals", - tags=["goals-management"], - summary="Create organizational goal", - description="Create a new goal for an organization with specified targets and deadlines", -) -async def create_goal( - organization_id: str = Path(..., description="Organization ID"), - goal_data: GoalCreateRequest = Body(..., description="Goal creation data"), - created_by: Optional[str] = Query( - None, description="ID of user/agent creating the goal" - ), -): - """Create a new organizational goal""" - try: - from .goals_management_service import GoalType - - goal_id = await app.state.goals_service.create_goal( - organization_id=organization_id, - title=goal_data.title, - description=goal_data.description, - goal_type=GoalType(goal_data.goal_type), - target_value=goal_data.target_value, - target_unit=goal_data.target_unit, - target_deadline=goal_data.target_deadline, - priority_level=goal_data.priority_level, - success_criteria=goal_data.success_criteria, - assigned_teams=goal_data.assigned_teams, - goal_owner_agent_id=goal_data.goal_owner_agent_id, - stakeholder_agents=goal_data.stakeholder_agents, - tags=goal_data.tags, - metadata=goal_data.metadata, - created_by=created_by, - ) - - return {"goal_id": goal_id, "status": "created"} - - except Exception as e: - logger.error(f"Error creating goal: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/organizations/{organization_id}/goals", - tags=["goals-management"], - summary="List organization goals", - description="Get all goals for an organization with optional filtering", -) -async def list_organization_goals( - organization_id: str = Path(..., description="Organization ID"), - status: Optional[List[str]] = Query(None, description="Filter by goal status"), - goal_type: Optional[List[str]] = Query(None, description="Filter by goal type"), - limit: int = Query( - 50, ge=1, le=100, description="Maximum number of goals to return" - ), -): - """List goals for an organization""" - try: - from .goals_management_service import GoalStatus, GoalType - - status_filter = [GoalStatus(s) for s in status] if status else None - type_filter = [GoalType(gt) for gt in goal_type] if goal_type else None - - goals = await app.state.goals_service.list_organization_goals( - organization_id=organization_id, - status_filter=status_filter, - goal_type_filter=type_filter, - limit=limit, - ) - - return { - "organization_id": organization_id, - "goals": [ - { - "id": goal.id, - "title": goal.title, - "description": goal.description, - "goal_type": goal.goal_type.value, - "status": goal.status.value, - "progress_percentage": float(goal.progress_percentage), - "target_value": ( - float(goal.target_value) if goal.target_value else None - ), - "target_unit": goal.target_unit, - "current_value": ( - float(goal.current_value) if goal.current_value else None - ), - "target_deadline": goal.target_deadline.isoformat(), - "priority_level": goal.priority_level, - "completion_confidence": float(goal.completion_confidence), - "created_at": goal.created_at.isoformat(), - "updated_at": goal.updated_at.isoformat(), - } - for goal in goals - ], - } - - except Exception as e: - logger.error(f"Error listing organization goals: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/goals/{goal_id}", - tags=["goals-management"], - summary="Get goal details", - description="Get detailed information about a specific goal", -) -async def get_goal(goal_id: str = Path(..., description="Goal ID")): - """Get goal details""" - try: - goal = await app.state.goals_service.get_goal(goal_id) - - if not goal: - raise HTTPException(status_code=404, detail="Goal not found") - - return { - "id": goal.id, - "organization_id": goal.organization_id, - "title": goal.title, - "description": goal.description, - "goal_type": goal.goal_type.value, - "status": goal.status.value, - "progress_percentage": float(goal.progress_percentage), - "target_value": float(goal.target_value) if goal.target_value else None, - "target_unit": goal.target_unit, - "current_value": float(goal.current_value) if goal.current_value else None, - "success_criteria": goal.success_criteria, - "start_date": goal.start_date.isoformat(), - "target_deadline": goal.target_deadline.isoformat(), - "actual_completion_date": ( - goal.actual_completion_date.isoformat() - if goal.actual_completion_date - else None - ), - "priority_level": goal.priority_level, - "completion_confidence": float(goal.completion_confidence), - "assigned_teams": goal.assigned_teams, - "goal_owner_agent_id": goal.goal_owner_agent_id, - "stakeholder_agents": goal.stakeholder_agents, - "tags": goal.tags, - "metadata": goal.metadata, - "created_by": goal.created_by, - "created_at": goal.created_at.isoformat(), - "updated_at": goal.updated_at.isoformat(), - } - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting goal {goal_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/goals/{goal_id}/overview", - tags=["goals-management"], - summary="Get goal overview", - description="Get comprehensive overview of goal with milestones, tasks, and progress", -) -async def get_goal_overview(goal_id: str = Path(..., description="Goal ID")): - """Get comprehensive goal overview""" - try: - overview = await app.state.goals_service.get_goal_overview(goal_id) - - if not overview: - raise HTTPException(status_code=404, detail="Goal not found") - - return overview - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting goal overview {goal_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.put( - "/goals/{goal_id}/progress", - tags=["goals-management"], - summary="Update goal progress", - description="Update progress for a specific goal", -) -async def update_goal_progress( - goal_id: str = Path(..., description="Goal ID"), - progress_data: GoalUpdateRequest = Body(..., description="Progress update data"), - recorded_by: Optional[str] = Query( - None, description="ID of user/agent recording progress" - ), -): - """Update goal progress""" - try: - success = await app.state.goals_service.update_goal_progress( - goal_id=goal_id, - progress_percentage=progress_data.progress_percentage, - current_value=progress_data.current_value, - completion_confidence=progress_data.completion_confidence, - progress_notes=progress_data.notes, - recorded_by=recorded_by, - ) - - if not success: - raise HTTPException( - status_code=404, detail="Goal not found or no changes made" - ) - - return {"goal_id": goal_id, "status": "updated"} - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error updating goal progress {goal_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/goals/{goal_id}/milestones", - tags=["goals-management"], - summary="Create milestone", - description="Create a new milestone for a goal", -) -async def create_milestone( - goal_id: str = Path(..., description="Goal ID"), - milestone_data: MilestoneCreateRequest = Body( - ..., description="Milestone creation data" - ), - created_by: Optional[str] = Query( - None, description="ID of user/agent creating milestone" - ), -): - """Create milestone for goal""" - try: - milestone_id = await app.state.goals_service.create_milestone( - goal_id=goal_id, - title=milestone_data.title, - description=milestone_data.description, - target_date=milestone_data.target_date, - milestone_type=milestone_data.milestone_type, - success_criteria=milestone_data.success_criteria, - deliverables=milestone_data.deliverables, - dependencies=milestone_data.dependencies, - assigned_teams=milestone_data.assigned_teams, - responsible_agent_id=milestone_data.responsible_agent_id, - priority_level=milestone_data.priority_level, - weight_in_goal=milestone_data.weight_in_goal, - created_by=created_by, - ) - - return {"milestone_id": milestone_id, "status": "created"} - - except Exception as e: - logger.error(f"Error creating milestone: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/milestones/{milestone_id}/tasks", - tags=["goals-management"], - summary="Create task from milestone", - description="Create a new task derived from a milestone", -) -async def create_task_from_milestone( - milestone_id: str = Path(..., description="Milestone ID"), - task_data: TaskFromMilestoneRequest = Body(..., description="Task creation data"), - created_by: Optional[str] = Query( - None, description="ID of user/agent creating task" - ), -): - """Create task from milestone""" - try: - task_id = await app.state.goals_service.create_task_from_milestone( - milestone_id=milestone_id, - title=task_data.title, - description=task_data.description, - task_type=task_data.task_type, - complexity_level=task_data.complexity_level, - estimated_hours=task_data.estimated_hours, - due_date=task_data.due_date, - assigned_team_id=task_data.assigned_team_id, - assigned_agent_id=task_data.assigned_agent_id, - priority=task_data.priority, - requirements=task_data.requirements, - acceptance_criteria=task_data.acceptance_criteria, - dependencies=task_data.dependencies, - created_by_agent_id=created_by, - ) - - return {"task_id": task_id, "status": "created"} - - except Exception as e: - logger.error(f"Error creating task from milestone: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/goals/{goal_id}/generate-execution-plan", - tags=["goals-management"], - summary="Generate execution plan", - description="Generate comprehensive milestone and task execution plan for a goal", -) -async def generate_execution_plan( - goal_id: str = Path(..., description="Goal ID"), - planning_context: Optional[Dict[str, Any]] = Body( - None, description="Additional planning context" - ), -): - """Generate execution plan with milestones and tasks""" - try: - execution_plan = ( - await app.state.milestone_task_engine.generate_goal_execution_plan( - goal_id=goal_id, planning_context=planning_context - ) - ) - - return execution_plan - - except Exception as e: - logger.error(f"Error generating execution plan for goal {goal_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/goals/{goal_id}/generate-monthly-milestones", - tags=["goals-management"], - summary="Generate monthly milestones", - description="Generate monthly milestone breakdown for a goal", -) -async def generate_monthly_milestones( - goal_id: str = Path(..., description="Goal ID"), - start_date: Optional[date] = Query(None, description="Start date for milestones"), - end_date: Optional[date] = Query(None, description="End date for milestones"), -): - """Generate monthly milestones for goal""" - try: - milestone_ids = ( - await app.state.milestone_task_engine.generate_monthly_milestones( - goal_id=goal_id, start_date=start_date, end_date=end_date - ) - ) - - return { - "goal_id": goal_id, - "milestone_ids": milestone_ids, - "count": len(milestone_ids), - "status": "generated", - } - - except Exception as e: - logger.error(f"Error generating monthly milestones for goal {goal_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/milestones/{milestone_id}/generate-weekly-tasks", - tags=["goals-management"], - summary="Generate weekly tasks", - description="Generate weekly task breakdown for a milestone", -) -async def generate_weekly_tasks( - milestone_id: str = Path(..., description="Milestone ID"), - focus_areas: Optional[List[str]] = Body( - None, description="Focus areas for task generation" - ), -): - """Generate weekly tasks for milestone""" - try: - task_ids = ( - await app.state.milestone_task_engine.generate_weekly_tasks_for_milestone( - milestone_id=milestone_id, focus_areas=focus_areas - ) - ) - - return { - "milestone_id": milestone_id, - "task_ids": task_ids, - "count": len(task_ids), - "status": "generated", - } - - except Exception as e: - logger.error(f"Error generating weekly tasks for milestone {milestone_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/goals/{goal_id}/generate-cross-functional-tasks", - tags=["goals-management"], - summary="Generate cross-functional tasks", - description="Generate tasks across different business functions for a goal", -) -async def generate_cross_functional_tasks( - goal_id: str = Path(..., description="Goal ID"), - target_functions: Optional[List[str]] = Body( - None, description="Target business functions" - ), -): - """Generate cross-functional tasks for goal""" - try: - functional_tasks = ( - await app.state.milestone_task_engine.generate_cross_functional_tasks( - goal_id=goal_id, target_functions=target_functions - ) - ) - - return { - "goal_id": goal_id, - "functional_tasks": functional_tasks, - "total_tasks": sum(len(tasks) for tasks in functional_tasks.values()), - "status": "generated", - } - - except Exception as e: - logger.error(f"Error generating cross-functional tasks for goal {goal_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/goals/{goal_id}/conversations", - tags=["goals-management"], - summary="Create goal conversation", - description="Create AI-powered conversation for goal planning and discussion", -) -async def create_goal_conversation( - goal_id: str = Path(..., description="Goal ID"), - conversation_data: GoalConversationCreateRequest = Body( - ..., description="Conversation creation data" - ), - created_by: Optional[str] = Query( - None, description="ID of user/agent creating conversation" - ), -): - """Create goal conversation""" - try: - from .goal_conversation_service import ConversationType - - conversation_id = ( - await app.state.goal_conversation_service.create_goal_conversation( - goal_id=goal_id, - conversation_type=ConversationType(conversation_data.conversation_type), - conversation_title=conversation_data.conversation_title, - initial_context=conversation_data.initial_context, - participants=conversation_data.participants, - created_by=created_by, - ) - ) - - return {"conversation_id": conversation_id, "status": "created"} - - except Exception as e: - logger.error(f"Error creating goal conversation: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/conversations/{conversation_id}", - tags=["goals-management"], - summary="Get goal conversation", - description="Get full conversation with messages, insights, and action items", -) -async def get_goal_conversation( - conversation_id: str = Path(..., description="Conversation ID") -): - """Get goal conversation""" - try: - conversation = await app.state.goal_conversation_service.get_conversation( - conversation_id - ) - - if not conversation: - raise HTTPException(status_code=404, detail="Conversation not found") - - return conversation - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting conversation {conversation_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/conversations/{conversation_id}/messages", - tags=["goals-management"], - summary="Add message to conversation", - description="Add a new message to a goal conversation", -) -async def add_message_to_conversation( - conversation_id: str = Path(..., description="Conversation ID"), - message_data: ConversationMessageRequest = Body(..., description="Message data"), - sender_id: Optional[str] = Query(None, description="ID of message sender"), -): - """Add message to conversation""" - try: - from .goal_conversation_service import MessageType - - message_id = ( - await app.state.goal_conversation_service.add_message_to_conversation( - conversation_id=conversation_id, - message_type=MessageType(message_data.message_type), - sender_id=sender_id, - sender_name=message_data.sender_name, - content=message_data.content, - metadata=message_data.metadata, - references=message_data.references, - ) - ) - - return {"message_id": message_id, "status": "added"} - - except Exception as e: - logger.error(f"Error adding message to conversation: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/conversations/{conversation_id}/generate-milestones", - tags=["goals-management"], - summary="Generate milestones from conversation", - description="Generate milestone recommendations based on conversation analysis", -) -async def generate_planning_milestones( - conversation_id: str = Path(..., description="Conversation ID"), - planning_context: Optional[Dict[str, Any]] = Body( - None, description="Additional planning context" - ), -): - """Generate planning milestones from conversation""" - try: - milestones = ( - await app.state.goal_conversation_service.generate_planning_milestones( - conversation_id=conversation_id, planning_context=planning_context - ) - ) - - return { - "conversation_id": conversation_id, - "milestones": milestones, - "count": len(milestones), - "status": "generated", - } - - except Exception as e: - logger.error(f"Error generating planning milestones: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/conversations/{conversation_id}/conduct-progress-review", - tags=["goals-management"], - summary="Conduct progress review", - description="Conduct AI-powered progress review for a goal conversation", -) -async def conduct_progress_review( - conversation_id: str = Path(..., description="Conversation ID"), - review_period_days: int = Query( - 30, ge=1, le=365, description="Review period in days" - ), -): - """Conduct progress review""" - try: - review_analysis = ( - await app.state.goal_conversation_service.conduct_progress_review( - conversation_id=conversation_id, review_period_days=review_period_days - ) - ) - - return review_analysis - - except Exception as e: - logger.error(f"Error conducting progress review: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/conversations/{conversation_id}/extract-action-items", - tags=["goals-management"], - summary="Extract action items", - description="Extract and create action items from conversation analysis", -) -async def extract_action_items( - conversation_id: str = Path(..., description="Conversation ID"), - auto_assign: bool = Query( - True, description="Whether to automatically assign action items" - ), -): - """Extract action items from conversation""" - try: - action_items = await app.state.goal_conversation_service.extract_action_items_from_conversation( - conversation_id=conversation_id, auto_assign=auto_assign - ) - - return { - "conversation_id": conversation_id, - "action_items": action_items, - "count": len(action_items), - "status": "extracted", - } - - except Exception as e: - logger.error(f"Error extracting action items: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/goals/{goal_id}/conversations", - tags=["goals-management"], - summary="Get goal conversations", - description="Get all conversations for a goal with optional filtering", -) -async def get_goal_conversations( - goal_id: str = Path(..., description="Goal ID"), - conversation_type: Optional[str] = Query( - None, description="Filter by conversation type" - ), - status: Optional[str] = Query(None, description="Filter by conversation status"), - limit: int = Query(10, ge=1, le=50, description="Maximum number of conversations"), -): - """Get conversations for a goal""" - try: - from .goal_conversation_service import ConversationStatus, ConversationType - - conv_type = ConversationType(conversation_type) if conversation_type else None - conv_status = ConversationStatus(status) if status else None - - conversations = ( - await app.state.goal_conversation_service.get_goal_conversations( - goal_id=goal_id, - conversation_type=conv_type, - status=conv_status, - limit=limit, - ) - ) - - return { - "goal_id": goal_id, - "conversations": conversations, - "count": len(conversations), - } - - except Exception as e: - logger.error(f"Error getting goal conversations: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/goals/{goal_id}/track-progress", - tags=["goals-management"], - summary="Record progress tracking update", - description="Record detailed progress update with tracking and risk assessment", -) -async def record_progress_tracking( - goal_id: str = Path(..., description="Goal ID"), - progress_data: ProgressUpdateRequest = Body( - ..., description="Progress tracking data" - ), - recorded_by: Optional[str] = Query( - None, description="ID of user/agent recording progress" - ), -): - """Record progress tracking update""" - try: - snapshot_id = await app.state.goal_tracking_service.record_progress_update( - goal_id=goal_id, - progress_percentage=progress_data.progress_percentage, - current_value=progress_data.current_value, - milestone_id=progress_data.milestone_id, - notes=progress_data.notes, - recorded_by=recorded_by, - confidence_score=progress_data.confidence_score, - trigger_alerts=progress_data.trigger_alerts, - ) - - return {"goal_id": goal_id, "snapshot_id": snapshot_id, "status": "recorded"} - - except Exception as e: - logger.error(f"Error recording progress tracking: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/goals/{goal_id}/deadline-risk", - tags=["goals-management"], - summary="Assess deadline risk", - description="Get comprehensive deadline risk assessment for a goal", -) -async def assess_deadline_risk(goal_id: str = Path(..., description="Goal ID")): - """Assess deadline risk for goal""" - try: - deadline_risk = await app.state.goal_tracking_service.assess_goal_deadline_risk( - goal_id - ) - - return { - "goal_id": deadline_risk.goal_id, - "risk_level": deadline_risk.risk_level.value, - "probability_of_delay": float(deadline_risk.probability_of_delay), - "estimated_completion_date": deadline_risk.estimated_completion_date.isoformat(), - "days_at_risk": deadline_risk.days_at_risk, - "critical_path_items": deadline_risk.critical_path_items, - "mitigation_strategies": deadline_risk.mitigation_strategies, - "updated_at": deadline_risk.updated_at.isoformat(), - } - - except Exception as e: - logger.error(f"Error assessing deadline risk: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/goals/{goal_id}/progress-report", - tags=["goals-management"], - summary="Generate progress report", - description="Generate comprehensive progress report for a goal", -) -async def generate_progress_report( - goal_id: str = Path(..., description="Goal ID"), - report_period_days: int = Query( - 30, ge=1, le=365, description="Report period in days" - ), -): - """Generate progress report for goal""" - try: - report = await app.state.goal_tracking_service.generate_progress_report( - goal_id=goal_id, report_period_days=report_period_days - ) - - return report - - except Exception as e: - logger.error(f"Error generating progress report: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/organizations/{organization_id}/goals-dashboard", - tags=["goals-management"], - summary="Get organization goals dashboard", - description="Get comprehensive dashboard for all organization goals", -) -async def get_organization_goals_dashboard( - organization_id: str = Path(..., description="Organization ID") -): - """Get organization goals dashboard""" - try: - dashboard = await app.state.goals_service.get_organization_goals_dashboard( - organization_id - ) - return dashboard - - except Exception as e: - logger.error(f"Error getting organization dashboard: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/organizations/{organization_id}/tracking-dashboard", - tags=["goals-management"], - summary="Get tracking dashboard", - description="Get comprehensive tracking dashboard with risk assessments", -) -async def get_tracking_dashboard( - organization_id: str = Path(..., description="Organization ID") -): - """Get organization tracking dashboard""" - try: - dashboard = ( - await app.state.goal_tracking_service.get_organization_tracking_dashboard( - organization_id - ) - ) - return dashboard - - except Exception as e: - logger.error(f"Error getting tracking dashboard: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# ================================ -# Knowledge Management API Endpoints -# ================================ - - -@app.post( - "/knowledge/organizations/{organization_id}/documents", - tags=["knowledge-management"], - summary="Upload Organizational Document", - response_model=DocumentMetadata, -) -async def upload_organization_document( - organization_id: str = Path(..., description="Organization ID"), - file: UploadFile = File(..., description="Document file to upload"), - title: Optional[str] = Form(None, description="Document title"), - tags: Optional[str] = Form(None, description="Comma-separated tags"), -): - """Upload a document to organizational knowledge base""" - try: - tags_list = [] - if tags: - tags_list = [tag.strip() for tag in tags.split(",")] - - document = await knowledge_manager.upload_document( - file_content=file.file, - filename=file.filename, - title=title, - organization_id=organization_id, - tags=tags_list, - ) - - return document - - except Exception as e: - logger.error(f"Error uploading organizational document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/knowledge/organizations/{organization_id}/url", - tags=["knowledge-management"], - summary="Add URL to Organizational Knowledge", - response_model=DocumentMetadata, -) -async def add_organization_url( - organization_id: str = Path(..., description="Organization ID"), - url: str = Body(..., embed=True), - title: Optional[str] = Body(None, embed=True), - tags: Optional[List[str]] = Body(None, embed=True), -): - """Add URL content to organizational knowledge base""" - try: - document = await knowledge_manager.upload_url( - url=url, title=title, organization_id=organization_id, tags=tags or [] - ) - - return document - - except Exception as e: - logger.error(f"Error adding organizational URL: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/organizations/{organization_id}/documents", - tags=["knowledge-management"], - summary="List Organizational Documents", - response_model=List[DocumentMetadata], -) -async def list_organization_documents( - organization_id: str = Path(..., description="Organization ID") -): - """Get list of organizational documents""" - try: - documents = await knowledge_manager.get_documents( - organization_id=organization_id - ) - return documents - - except Exception as e: - logger.error(f"Error listing organizational documents: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/organizations/{organization_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Get Organizational Document", - response_model=DocumentMetadata, -) -async def get_organization_document( - organization_id: str = Path(..., description="Organization ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Get organizational document metadata""" - try: - document = await knowledge_manager.get_document_metadata( - doc_id=doc_id, organization_id=organization_id - ) - - if not document: - raise HTTPException(status_code=404, detail="Document not found") - - return document - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting organizational document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/organizations/{organization_id}/documents/{doc_id}/content", - tags=["knowledge-management"], - summary="Get Organizational Document Content", -) -async def get_organization_document_content( - organization_id: str = Path(..., description="Organization ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Get full content of organizational document""" - try: - content = await knowledge_manager.get_document_content( - doc_id=doc_id, organization_id=organization_id - ) - - if content is None: - raise HTTPException(status_code=404, detail="Document not found") - - return {"content": content} - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting organizational document content: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.put( - "/knowledge/organizations/{organization_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Update Organizational Document", - response_model=DocumentMetadata, -) -async def update_organization_document( - organization_id: str = Path(..., description="Organization ID"), - doc_id: str = Path(..., description="Document ID"), - title: Optional[str] = Body(None, embed=True), - tags: Optional[List[str]] = Body(None, embed=True), -): - """Update organizational document metadata""" - try: - document = await knowledge_manager.update_document( - doc_id=doc_id, title=title, tags=tags, organization_id=organization_id - ) - - if not document: - raise HTTPException(status_code=404, detail="Document not found") - - return document - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error updating organizational document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.delete( - "/knowledge/organizations/{organization_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Delete Organizational Document", -) -async def delete_organization_document( - organization_id: str = Path(..., description="Organization ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Delete organizational document""" - try: - success = await knowledge_manager.delete_document( - doc_id=doc_id, organization_id=organization_id - ) - - if not success: - raise HTTPException(status_code=404, detail="Document not found") - - return {"message": "Document deleted successfully"} - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error deleting organizational document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# Team Knowledge Management Endpoints - - -@app.post( - "/knowledge/teams/{team_id}/documents", - tags=["knowledge-management"], - summary="Upload Team Document", - response_model=DocumentMetadata, -) -async def upload_team_document( - team_id: str = Path(..., description="Team ID"), - file: UploadFile = File(..., description="Document file to upload"), - title: Optional[str] = Form(None, description="Document title"), - tags: Optional[str] = Form(None, description="Comma-separated tags"), -): - """Upload a document to team knowledge base""" - try: - tags_list = [] - if tags: - tags_list = [tag.strip() for tag in tags.split(",")] - - document = await knowledge_manager.upload_document( - file_content=file.file, - filename=file.filename, - title=title, - team_id=team_id, - tags=tags_list, - ) - - return document - - except Exception as e: - logger.error(f"Error uploading team document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/knowledge/teams/{team_id}/url", - tags=["knowledge-management"], - summary="Add URL to Team Knowledge", - response_model=DocumentMetadata, -) -async def add_team_url( - team_id: str = Path(..., description="Team ID"), - url: str = Body(..., embed=True), - title: Optional[str] = Body(None, embed=True), - tags: Optional[List[str]] = Body(None, embed=True), -): - """Add URL content to team knowledge base""" - try: - document = await knowledge_manager.upload_url( - url=url, title=title, team_id=team_id, tags=tags or [] - ) - - return document - - except Exception as e: - logger.error(f"Error adding team URL: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/teams/{team_id}/documents", - tags=["knowledge-management"], - summary="List Team Documents", - response_model=List[DocumentMetadata], -) -async def list_team_documents(team_id: str = Path(..., description="Team ID")): - """Get list of team documents""" - try: - documents = await knowledge_manager.get_documents(team_id=team_id) - return documents - - except Exception as e: - logger.error(f"Error listing team documents: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/teams/{team_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Get Team Document", - response_model=DocumentMetadata, -) -async def get_team_document( - team_id: str = Path(..., description="Team ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Get team document metadata""" - try: - document = await knowledge_manager.get_document_metadata( - doc_id=doc_id, team_id=team_id - ) - - if not document: - raise HTTPException(status_code=404, detail="Document not found") - - return document - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting team document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/teams/{team_id}/documents/{doc_id}/content", - tags=["knowledge-management"], - summary="Get Team Document Content", -) -async def get_team_document_content( - team_id: str = Path(..., description="Team ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Get full content of team document""" - try: - content = await knowledge_manager.get_document_content( - doc_id=doc_id, team_id=team_id - ) - - if content is None: - raise HTTPException(status_code=404, detail="Document not found") - - return {"content": content} - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting team document content: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.put( - "/knowledge/teams/{team_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Update Team Document", - response_model=DocumentMetadata, -) -async def update_team_document( - team_id: str = Path(..., description="Team ID"), - doc_id: str = Path(..., description="Document ID"), - title: Optional[str] = Body(None, embed=True), - tags: Optional[List[str]] = Body(None, embed=True), -): - """Update team document metadata""" - try: - document = await knowledge_manager.update_document( - doc_id=doc_id, title=title, tags=tags, team_id=team_id - ) - - if not document: - raise HTTPException(status_code=404, detail="Document not found") - - return document - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error updating team document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.delete( - "/knowledge/teams/{team_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Delete Team Document", -) -async def delete_team_document( - team_id: str = Path(..., description="Team ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Delete team document""" - try: - success = await knowledge_manager.delete_document( - doc_id=doc_id, team_id=team_id - ) - - if not success: - raise HTTPException(status_code=404, detail="Document not found") - - return {"message": "Document deleted successfully"} - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error deleting team document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# Agent Knowledge Management Endpoints - - -@app.post( - "/knowledge/agents/{agent_id}/documents", - tags=["knowledge-management"], - summary="Upload Agent Document", - response_model=DocumentMetadata, -) -async def upload_agent_document( - agent_id: str = Path(..., description="Agent ID"), - file: UploadFile = File(..., description="Document file to upload"), - title: Optional[str] = Form(None, description="Document title"), - tags: Optional[str] = Form(None, description="Comma-separated tags"), -): - """Upload a document to agent knowledge base""" - try: - tags_list = [] - if tags: - tags_list = [tag.strip() for tag in tags.split(",")] - - document = await knowledge_manager.upload_document( - file_content=file.file, - filename=file.filename, - title=title, - agent_id=agent_id, - tags=tags_list, - ) - - return document - - except Exception as e: - logger.error(f"Error uploading agent document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/knowledge/agents/{agent_id}/url", - tags=["knowledge-management"], - summary="Add URL to Agent Knowledge", - response_model=DocumentMetadata, -) -async def add_agent_url( - agent_id: str = Path(..., description="Agent ID"), - url: str = Body(..., embed=True), - title: Optional[str] = Body(None, embed=True), - tags: Optional[List[str]] = Body(None, embed=True), -): - """Add URL content to agent knowledge base""" - try: - document = await knowledge_manager.upload_url( - url=url, title=title, agent_id=agent_id, tags=tags or [] - ) - - return document - - except Exception as e: - logger.error(f"Error adding agent URL: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/agents/{agent_id}/documents", - tags=["knowledge-management"], - summary="List Agent Documents", - response_model=List[DocumentMetadata], -) -async def list_agent_documents(agent_id: str = Path(..., description="Agent ID")): - """Get list of agent documents""" - try: - documents = await knowledge_manager.get_documents(agent_id=agent_id) - return documents - - except Exception as e: - logger.error(f"Error listing agent documents: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/agents/{agent_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Get Agent Document", - response_model=DocumentMetadata, -) -async def get_agent_document( - agent_id: str = Path(..., description="Agent ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Get agent document metadata""" - try: - document = await knowledge_manager.get_document_metadata( - doc_id=doc_id, agent_id=agent_id - ) - - if not document: - raise HTTPException(status_code=404, detail="Document not found") - - return document - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting agent document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/agents/{agent_id}/documents/{doc_id}/content", - tags=["knowledge-management"], - summary="Get Agent Document Content", -) -async def get_agent_document_content( - agent_id: str = Path(..., description="Agent ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Get full content of agent document""" - try: - content = await knowledge_manager.get_document_content( - doc_id=doc_id, agent_id=agent_id - ) - - if content is None: - raise HTTPException(status_code=404, detail="Document not found") - - return {"content": content} - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting agent document content: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.put( - "/knowledge/agents/{agent_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Update Agent Document", - response_model=DocumentMetadata, -) -async def update_agent_document( - agent_id: str = Path(..., description="Agent ID"), - doc_id: str = Path(..., description="Document ID"), - title: Optional[str] = Body(None, embed=True), - tags: Optional[List[str]] = Body(None, embed=True), -): - """Update agent document metadata""" - try: - document = await knowledge_manager.update_document( - doc_id=doc_id, title=title, tags=tags, agent_id=agent_id - ) - - if not document: - raise HTTPException(status_code=404, detail="Document not found") - - return document - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error updating agent document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.delete( - "/knowledge/agents/{agent_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Delete Agent Document", -) -async def delete_agent_document( - agent_id: str = Path(..., description="Agent ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Delete agent document""" - try: - success = await knowledge_manager.delete_document( - doc_id=doc_id, agent_id=agent_id - ) - - if not success: - raise HTTPException(status_code=404, detail="Document not found") - - return {"message": "Document deleted successfully"} - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error deleting agent document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# Knowledge Search Endpoints - - -@app.get( - "/knowledge/search", - tags=["knowledge-management"], - summary="Search Knowledge Base", - response_model=List[DocumentMetadata], -) -async def search_knowledge( - query: str = Query(..., description="Search query"), - organization_id: Optional[str] = Query(None, description="Filter by organization"), - team_id: Optional[str] = Query(None, description="Filter by team"), - agent_id: Optional[str] = Query(None, description="Filter by agent"), - limit: int = Query(10, ge=1, le=100, description="Maximum number of results"), -): - """Search across knowledge base""" - try: - documents = await knowledge_manager.search_documents( - query=query, - organization_id=organization_id, - team_id=team_id, - agent_id=agent_id, - limit=limit, - ) - - return documents - - except Exception as e: - logger.error(f"Error searching knowledge: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# ================================ -# Container Management API Endpoints -# ================================ - - -@app.post( - "/agents/{agent_id}/container/create", - tags=["container-management"], - summary="Create Agent Container", - response_model=ContainerStatus, -) -async def create_agent_container( - agent_id: str = Path(..., description="Agent ID"), - config: Optional[ContainerConfig] = Body( - None, description="Container configuration" - ), -): - """Create a new container for an AI agent""" - try: - status = await container_manager.create_agent_container(agent_id, config) - return status - - except Exception as e: - logger.error(f"Error creating container for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/agents/{agent_id}/container/start", - tags=["container-management"], - summary="Start Agent Container", - response_model=ContainerStatus, -) -async def start_agent_container(agent_id: str = Path(..., description="Agent ID")): - """Start an agent container""" - try: - status = await container_manager.start_container(agent_id) - return status - - except RuntimeError as e: - raise HTTPException(status_code=404, detail=str(e)) - except Exception as e: - logger.error(f"Error starting container for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/agents/{agent_id}/container/stop", - tags=["container-management"], - summary="Stop Agent Container", - response_model=ContainerStatus, -) -async def stop_agent_container( - agent_id: str = Path(..., description="Agent ID"), - timeout: int = Body(30, description="Stop timeout in seconds"), -): - """Stop an agent container""" - try: - status = await container_manager.stop_container(agent_id, timeout) - return status - - except RuntimeError as e: - raise HTTPException(status_code=404, detail=str(e)) - except Exception as e: - logger.error(f"Error stopping container for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/agents/{agent_id}/container/restart", - tags=["container-management"], - summary="Restart Agent Container", - response_model=ContainerStatus, -) -async def restart_agent_container( - agent_id: str = Path(..., description="Agent ID"), - timeout: int = Body(30, description="Restart timeout in seconds"), -): - """Restart an agent container""" - try: - status = await container_manager.restart_container(agent_id, timeout) - return status - - except RuntimeError as e: - raise HTTPException(status_code=404, detail=str(e)) - except Exception as e: - logger.error(f"Error restarting container for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.delete( - "/agents/{agent_id}/container", - tags=["container-management"], - summary="Remove Agent Container", -) -async def remove_agent_container( - agent_id: str = Path(..., description="Agent ID"), - force: bool = Query(False, description="Force removal of running container"), -): - """Remove an agent container""" - try: - success = await container_manager.remove_container(agent_id, force) - - if success: - return {"message": f"Container for agent {agent_id} removed successfully"} - else: - raise HTTPException(status_code=500, detail="Failed to remove container") - - except Exception as e: - logger.error(f"Error removing container for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/agents/{agent_id}/container/status", - tags=["container-management"], - summary="Get Agent Container Status", - response_model=Optional[ContainerStatus], -) -async def get_agent_container_status(agent_id: str = Path(..., description="Agent ID")): - """Get container status for an agent""" - try: - status = await container_manager.get_container_status(agent_id) - return status - - except Exception as e: - logger.error(f"Error getting container status for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/containers/agents", - tags=["container-management"], - summary="List Agent Containers", - response_model=List[ContainerStatus], -) -async def list_agent_containers(): - """List all agent containers""" - try: - containers = await container_manager.list_agent_containers() - return containers - - except Exception as e: - logger.error(f"Error listing agent containers: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/agents/{agent_id}/container/logs", - tags=["container-management"], - summary="Get Agent Container Logs", -) -async def get_agent_container_logs( - agent_id: str = Path(..., description="Agent ID"), - tail: int = Query(100, ge=1, le=10000, description="Number of log lines to return"), - since: Optional[str] = Query( - None, description="Show logs since timestamp (ISO format)" - ), -): - """Get container logs for an agent""" - try: - since_dt = None - if since: - try: - since_dt = datetime.fromisoformat(since.replace("Z", "+00:00")) - except ValueError: - raise HTTPException(status_code=400, detail="Invalid timestamp format") - - logs = await container_manager.get_container_logs( - agent_id=agent_id, tail=tail, since=since_dt - ) - - return {"logs": logs} - - except RuntimeError as e: - raise HTTPException(status_code=404, detail=str(e)) - except Exception as e: - logger.error(f"Error getting container logs for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/agents/{agent_id}/container/execute", - tags=["container-management"], - summary="Execute Command in Container", -) -async def execute_container_command( - agent_id: str = Path(..., description="Agent ID"), - command: str = Body(..., description="Command to execute"), - working_dir: Optional[str] = Body(None, description="Working directory"), -): - """Execute a command in the agent container""" - try: - result = await container_manager.execute_command( - agent_id=agent_id, command=command, working_dir=working_dir - ) - - return result - - except RuntimeError as e: - raise HTTPException(status_code=404, detail=str(e)) - except Exception as e: - logger.error(f"Error executing command in container for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# WebSocket endpoint for real-time log streaming -@app.websocket("/agents/{agent_id}/container/logs/stream") -async def stream_agent_container_logs(websocket: WebSocket, agent_id: str): - """Stream container logs in real-time via WebSocket""" - await websocket.accept() - - try: - # Check if container exists - status = await container_manager.get_container_status(agent_id) - if not status: - await websocket.send_json({"error": "Container not found"}) - await websocket.close() - return - - await websocket.send_json({"status": "connected", "agent_id": agent_id}) - - # Stream logs - async for log_entry in container_manager.stream_container_logs(agent_id): - await websocket.send_json( - { - "timestamp": log_entry.timestamp.isoformat(), - "stream": log_entry.stream, - "message": log_entry.message, - } - ) - - except Exception as e: - logger.error(f"Error in log stream for agent {agent_id}: {e}") - try: - await websocket.send_json({"error": str(e)}) - except Exception: - logger.debug("Failed to send error frame on closing websocket") - finally: - try: - await websocket.close() - except Exception: - logger.debug("Failed to close websocket cleanly") - - -# ============================================================================ -# RAG (Retrieval-Augmented Generation) Endpoints -# ============================================================================ - - -@app.post("/rag/search") -async def search_knowledge_context( - query: str = Body(..., embed=True), - organization_id: Optional[str] = Body(None, embed=True), - team_id: Optional[str] = Body(None, embed=True), - agent_id: Optional[str] = Body(None, embed=True), - max_results: int = Body(5, embed=True), - similarity_threshold: float = Body(0.7, embed=True), -): - """Search for relevant knowledge context using RAG""" - try: - context = await rag_system.search_relevant_context( - query=query, - organization_id=organization_id, - team_id=team_id, - agent_id=agent_id, - max_results=max_results, - similarity_threshold=similarity_threshold, - ) - - return { - "query": context.query, - "relevant_chunks": [ - { - "document_id": chunk.document_id, - "document_title": chunk.metadata.get("document_title", "Unknown"), - "content": chunk.content, - "chunk_index": chunk.chunk_index, - "metadata": chunk.metadata, - } - for chunk in context.relevant_chunks - ], - "similarity_scores": context.similarity_scores, - "total_documents": context.total_documents, - "context_length": context.context_length, - } - - except Exception as e: - logger.error(f"Error searching knowledge context: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post("/rag/enhance-prompt") -async def enhance_prompt_with_context( - message: str = Body(..., embed=True), - organization_id: Optional[str] = Body(None, embed=True), - team_id: Optional[str] = Body(None, embed=True), - agent_id: Optional[str] = Body(None, embed=True), - max_context_length: int = Body(4000, embed=True), -): - """Enhance a prompt with relevant context using RAG""" - try: - enhanced_prompt = await rag_system.get_contextual_prompt( - user_message=message, - organization_id=organization_id, - team_id=team_id, - agent_id=agent_id, - max_context_length=max_context_length, - ) - - return { - "original_message": message, - "enhanced_prompt": enhanced_prompt, - "context_added": len(enhanced_prompt) > len(message), - } - - except Exception as e: - logger.error(f"Error enhancing prompt with context: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post("/rag/reindex") -async def reindex_knowledge_base( - organization_id: Optional[str] = Body(None, embed=True), - team_id: Optional[str] = Body(None, embed=True), - agent_id: Optional[str] = Body(None, embed=True), -): - """Reindex all documents in a scope for RAG""" - try: - results = await rag_system.index_all_documents( - organization_id=organization_id, team_id=team_id, agent_id=agent_id - ) - - return { - "scope": { - "organization_id": organization_id, - "team_id": team_id, - "agent_id": agent_id, - }, - "results": results, - "message": f"Indexed {results['indexed']} documents, {results['failed']} failed, {results['skipped']} skipped", - } - - except Exception as e: - logger.error(f"Error reindexing knowledge base: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get("/rag/stats") -async def get_rag_index_stats(): - """Get statistics about the RAG index""" - try: - stats = await rag_system.get_index_stats() - return stats - - except Exception as e: - logger.error(f"Error getting RAG stats: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post("/rag/documents/{doc_id}/reindex") -async def reindex_document( - doc_id: str, - organization_id: Optional[str] = Body(None, embed=True), - team_id: Optional[str] = Body(None, embed=True), - agent_id: Optional[str] = Body(None, embed=True), -): - """Reindex a specific document for RAG""" - try: - # Get document metadata - document = await knowledge_manager.get_document_metadata( - doc_id=doc_id, - organization_id=organization_id, - team_id=team_id, - agent_id=agent_id, - ) - - if not document: - raise HTTPException(status_code=404, detail="Document not found") - - # Reindex the document - success = await rag_system.index_document(document) - - if success: - return { - "document_id": doc_id, - "status": "reindexed", - "message": f"Document '{document.title}' has been reindexed successfully", - } - else: - raise HTTPException(status_code=500, detail="Failed to reindex document") - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error reindexing document {doc_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# ============================================================================ -# Real-time WebSocket Endpoints -# ============================================================================ - - -@app.websocket("/ws/updates") -async def websocket_real_time_updates( - websocket: WebSocket, - organization_id: Optional[str] = None, - team_id: Optional[str] = None, - agent_id: Optional[str] = None, - user_id: Optional[str] = None, - subscriptions: Optional[str] = None, -): - """Main WebSocket endpoint for real-time updates""" - import uuid - - connection_id = str(uuid.uuid4()) - - # Parse subscriptions - subscription_list = [] - if subscriptions: - subscription_list = subscriptions.split(",") - - try: - connection = await websocket_manager.connect( - websocket=websocket, - connection_id=connection_id, - organization_id=organization_id, - team_id=team_id, - agent_id=agent_id, - user_id=user_id, - subscriptions=subscription_list, - ) - - # Keep connection alive and handle pings - while True: - try: - # Wait for ping messages or disconnection - message = await websocket.receive_text() - - # Handle ping/pong - if message == "ping": - await websocket.send_text("pong") - connection.last_ping = datetime.now() - else: - # Parse other messages (subscription updates, etc.) - try: - data = json.loads(message) - if data.get("type") == "subscribe": - # Update subscriptions - new_subs = data.get("subscriptions", []) - connection.scope.subscriptions.clear() - for sub in new_subs: - try: - connection.scope.subscriptions.add(UpdateType(sub)) - except ValueError: - pass - - await connection.send_update( - WebSocketUpdate( - type=UpdateType.SYSTEM_NOTIFICATION, - data={ - "message": "Subscriptions updated", - "subscriptions": list( - connection.scope.subscriptions - ), - }, - ) - ) - except json.JSONDecodeError: - pass - - except WebSocketDisconnect: - break - - except Exception as e: - logger.error(f"WebSocket error for connection {connection_id}: {e}") - finally: - await websocket_manager.disconnect(connection_id) - - -@app.websocket("/ws/agent/{agent_id}/updates") -async def websocket_agent_updates(websocket: WebSocket, agent_id: str): - """WebSocket endpoint for specific agent updates""" - import uuid - - connection_id = f"agent-{agent_id}-{uuid.uuid4()}" - - try: - connection = await websocket_manager.connect( - websocket=websocket, - connection_id=connection_id, - agent_id=agent_id, - subscriptions=[ - UpdateType.AGENT_STATUS.value, - UpdateType.TASK_STATUS.value, - UpdateType.TASK_PROGRESS.value, - UpdateType.CONTAINER_STATUS.value, - UpdateType.CHAT_MESSAGE.value, - UpdateType.CHAT_TYPING.value, - ], - ) - - # Keep connection alive - while True: - try: - message = await websocket.receive_text() - if message == "ping": - await websocket.send_text("pong") - connection.last_ping = datetime.now() - except WebSocketDisconnect: - break - - except Exception as e: - logger.error(f"Agent WebSocket error for {agent_id}: {e}") - finally: - await websocket_manager.disconnect(connection_id) - - -@app.websocket("/ws/agents/{agent_id}/conversations/{conversation_id}") -async def websocket_agent_conversation( - websocket: WebSocket, agent_id: str, conversation_id: str -): - """WebSocket endpoint for real-time agent conversation""" - import uuid - - connection_id = f"conversation-{conversation_id}-{uuid.uuid4()}" - - await websocket.accept() - - try: - # Store connection for broadcasting - active_conversations = getattr(app.state, "active_conversations", {}) - if conversation_id not in active_conversations: - active_conversations[conversation_id] = [] - active_conversations[conversation_id].append(websocket) - app.state.active_conversations = active_conversations - - while True: - try: - # Receive message from client - data = await websocket.receive_json() - - if data.get("type") == "ping": - await websocket.send_json({"type": "pong"}) - elif data.get("type") == "message": - # Handle new message - message_content = data.get("content", "") - if message_content: - # Store message in database - async with get_db_connection() as conn: - message_id = await conn.fetchval( - """ - INSERT INTO agent_conversations (session_id, agent_id, message_type, content) - VALUES ($1, $2, 'user', $3) - RETURNING id - """, - conversation_id, - agent_id, - message_content, - ) - - # Update session activity - await conn.execute( - """ - UPDATE chat_sessions - SET last_activity = CURRENT_TIMESTAMP, message_count = message_count + 1 - WHERE id = $1 - """, - conversation_id, - ) - - # Broadcast to all connected clients for this conversation - message_data = { - "type": "new_message", - "message": { - "id": str(message_id), - "conversation_id": conversation_id, - "role": "user", - "content": message_content, - "timestamp": datetime.now().isoformat(), - "status": "sent", - }, - } - - for conn in active_conversations.get(conversation_id, []): - try: - await conn.send_json(message_data) - except Exception: - # Connection might be closed; skip this subscriber - logger.debug("Skipped broadcast to a closed websocket") - - # TODO: Here we would trigger agent response generation - # For now, send a simple acknowledgment after a delay - await asyncio.sleep(1) - - agent_response = { - "type": "new_message", - "message": { - "id": str(uuid.uuid4()), - "conversation_id": conversation_id, - "role": "agent", - "content": f"I received your message: {message_content}", - "timestamp": datetime.now().isoformat(), - "status": "received", - }, - } - - for conn in active_conversations.get(conversation_id, []): - try: - await conn.send_json(agent_response) - except Exception: - # Connection might be closed; skip this subscriber - logger.debug("Skipped broadcast to a closed websocket") - - # Store agent response in database - async with get_db_connection() as conn: - await conn.execute( - """ - INSERT INTO agent_conversations (session_id, agent_id, message_type, content) - VALUES ($1, $2, 'agent', $3) - """, - conversation_id, - agent_id, - agent_response["message"]["content"], - ) - - except WebSocketDisconnect: - break - except Exception as e: - logger.error(f"Error in conversation WebSocket: {e}") - - except Exception as e: - logger.error(f"Conversation WebSocket error for {conversation_id}: {e}") - finally: - # Clean up connection - if ( - hasattr(app.state, "active_conversations") - and conversation_id in app.state.active_conversations - ): - if websocket in app.state.active_conversations[conversation_id]: - app.state.active_conversations[conversation_id].remove(websocket) - - -@app.websocket("/ws/organization/{organization_id}/updates") -async def websocket_organization_updates(websocket: WebSocket, organization_id: str): - """WebSocket endpoint for organization-wide updates""" - import uuid - - connection_id = f"org-{organization_id}-{uuid.uuid4()}" - - try: - connection = await websocket_manager.connect( - websocket=websocket, - connection_id=connection_id, - organization_id=organization_id, - subscriptions=[ - UpdateType.AGENT_CREATED.value, - UpdateType.AGENT_UPDATED.value, - UpdateType.AGENT_DELETED.value, - UpdateType.KNOWLEDGE_UPDATED.value, - UpdateType.KNOWLEDGE_INDEXED.value, - UpdateType.SYSTEM_NOTIFICATION.value, - ], - ) - - # Keep connection alive - while True: - try: - message = await websocket.receive_text() - if message == "ping": - await websocket.send_text("pong") - connection.last_ping = datetime.now() - except WebSocketDisconnect: - break - - except Exception as e: - logger.error(f"Organization WebSocket error for {organization_id}: {e}") - finally: - await websocket_manager.disconnect(connection_id) - - -# WebSocket Statistics Endpoint -@app.get("/ws/stats") -async def get_websocket_stats(): - """Get WebSocket connection statistics""" - try: - stats = websocket_manager.get_stats() - return stats - except Exception as e: - logger.error(f"Error getting WebSocket stats: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# Manual notification endpoints for testing -@app.post("/ws/test/agent/{agent_id}/status") -async def test_agent_status_notification( - agent_id: str, - status: str = Body(..., embed=True), - message: Optional[str] = Body(None, embed=True), -): - """Test endpoint to send agent status notifications""" - try: - await notify_agent_status_change( - agent_id=agent_id, - status=status, - additional_data={"message": message} if message else None, - ) - return {"status": "notification_sent", "agent_id": agent_id} - except Exception as e: - logger.error(f"Error sending test notification: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# ============================================================================ -# Missing API Endpoints (Goals, Teams, Organizations) -# ============================================================================ - - -@app.get("/teams") -async def get_teams(): - """Get list of teams""" - # Mock data for now - return [ - { - "id": "1", - "name": "Development Team", - "description": "Frontend and backend developers", - "member_count": 5, - "organization_id": "1", - }, - { - "id": "2", - "name": "Executive Team", - "description": "Leadership and strategy", - "member_count": 3, - "organization_id": "1", - }, - ] - - -@app.get("/organizations/{organization_id}/goals") -async def get_organization_goals(organization_id: str): - """Get goals for an organization""" - # Mock data for now - return [ - { - "id": "1", - "title": "Increase Development Velocity", - "description": "Improve team productivity and code quality", - "status": "active", - "progress": 75, - "organization_id": organization_id, - "created_at": "2024-01-15T10:00:00Z", - "due_date": "2024-12-31T23:59:59Z", - }, - { - "id": "2", - "title": "Enhance AI Capabilities", - "description": "Expand AI agent capabilities and intelligence", - "status": "active", - "progress": 50, - "organization_id": organization_id, - "created_at": "2024-02-01T10:00:00Z", - "due_date": "2024-11-30T23:59:59Z", - }, - ] - - -@app.get("/goals/{goal_id}") -async def get_goal_details(goal_id: str): - """Get detailed information about a specific goal""" - # Mock data for now - return { - "id": goal_id, - "title": "Increase Development Velocity", - "description": "Improve team productivity and code quality through better tooling, processes, and automation", - "status": "active", - "progress": 75, - "organization_id": "1", - "team_id": "1", - "created_at": "2024-01-15T10:00:00Z", - "updated_at": "2024-08-06T16:30:00Z", - "due_date": "2024-12-31T23:59:59Z", - "milestones": [ - { - "id": "1", - "title": "Implement CI/CD Pipeline", - "description": "Set up automated testing and deployment", - "status": "completed", - "progress": 100, - "due_date": "2024-03-15T23:59:59Z", - }, - { - "id": "2", - "title": "Enhance Code Review Process", - "description": "Streamline code review workflow with automated tools", - "status": "in_progress", - "progress": 80, - "due_date": "2024-09-30T23:59:59Z", - }, - { - "id": "3", - "title": "Deploy AI-Powered Testing", - "description": "Implement intelligent test generation and execution", - "status": "planned", - "progress": 25, - "due_date": "2024-12-15T23:59:59Z", - }, - ], - "metrics": { - "deployment_frequency": "Daily", - "lead_time": "2.3 days", - "mttr": "45 minutes", - "change_failure_rate": "5%", - }, - "assigned_agents": [ - {"id": "1", "name": "DevOps Agent", "role": "CI/CD Specialist"}, - {"id": "2", "name": "QA Agent", "role": "Test Automation Engineer"}, - ], - } - - -@app.get("/agents/{agent_id}/tasks") -async def get_agent_tasks_list(agent_id: str): - """Get tasks for a specific agent - GET method""" - try: - # Get tasks from task queue - tasks = await app.state.task_queue.get_agent_tasks(agent_id) - return {"agent_id": agent_id, "tasks": tasks} - except Exception as e: - logger.error(f"Error getting tasks for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) +import asyncio +import json +import logging +import os +from collections import defaultdict +from contextlib import asynccontextmanager +from datetime import date, datetime +from decimal import Decimal +from typing import Any, Dict, List, Optional + +import jwt +from fastapi import ( + Body, + Depends, + FastAPI, + File, + Form, + HTTPException, + Path, + Query, + UploadFile, + WebSocket, + WebSocketDisconnect, + status, +) +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import HTMLResponse, Response +from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer +from pydantic import BaseModel, Field + +from hierarchy_endpoints import router as hierarchy_router + +from .agent_manager import AgentManager +from .container_manager import ContainerConfig, ContainerStatus, container_manager +from .context_service import ContextService +from .database import get_db_connection +from .knowledge_manager import DocumentMetadata, knowledge_manager +from .rag_integration import RAGContext, rag_system +from .sandbox_manager import AgentSandboxManager +from .task_execution_engine import TaskExecutionEngine +from .task_queue import TaskQueue +from .websocket_manager import ( + UpdateType, + WebSocketUpdate, + notify_agent_status_change, + notify_container_status_change, + notify_knowledge_update, + notify_task_progress, + websocket_manager, +) + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Auth helpers (Track 3) +# --------------------------------------------------------------------------- +_security = HTTPBearer(auto_error=False) +_jwt_secret = os.environ.get("FUZEFRONT_JWT_SECRET", "") + + +def require_auth(credentials: HTTPAuthorizationCredentials = Depends(_security)): + """Verify FuzeFront JWT on mutating endpoints. Disabled when secret not set (dev).""" + if not _jwt_secret: + return None # Auth disabled when secret not configured (dev mode) + if not credentials: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing token" + ) + try: + payload = jwt.decode(credentials.credentials, _jwt_secret, algorithms=["HS256"]) + return payload + except jwt.ExpiredSignatureError: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, detail="Token expired" + ) + except jwt.InvalidTokenError: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token" + ) + + +# --------------------------------------------------------------------------- +# Agent relay state (Track 4) +# --------------------------------------------------------------------------- +# agent_id -> list of subscriber WebSockets watching that agent's session +agent_relay_subscribers: Dict[str, List[WebSocket]] = defaultdict(list) + + +# Pydantic models for API documentation +class AgentCreateRequest(BaseModel): + name: str = Field(..., description="Agent name") + role: str = Field(..., description="Agent role (e.g., 'Senior React Developer')") + type: str = Field(..., description="Agent type (e.g., 'developer', 'executive')") + config: Dict[str, Any] = Field( + default_factory=dict, description="Agent configuration" + ) + repository_settings: Dict[str, Any] = Field( + default_factory=dict, description="Repository settings" + ) + sandbox_settings: Dict[str, Any] = Field( + default_factory=dict, description="Sandbox settings" + ) + + +class TaskCreateRequest(BaseModel): + title: str = Field(..., description="Task title") + description: str = Field(..., description="Task description") + priority: str = Field( + default="medium", description="Task priority (low, medium, high)" + ) + metadata: Dict[str, Any] = Field( + default_factory=dict, description="Additional task metadata" + ) + + +class HumanResponseRequest(BaseModel): + response: str = Field(..., description="Human response to agent question") + + +class FileOperationApprovalRequest(BaseModel): + approved: bool = Field(..., description="Whether to approve the file operations") + reason: Optional[str] = Field( + None, description="Optional reason for approval/rejection" + ) + + +class ClaudeSessionInputRequest(BaseModel): + input: str = Field(..., description="Input to send to Claude SDK session") + + +class CoordinationRequest(BaseModel): + coordination_mode: str = Field( + default="collaborative", + description="Coordination mode (sequential, parallel, hierarchical, collaborative)", + ) + required_agents: Optional[List[str]] = Field( + None, description="Specific agents to include" + ) + required_skills: Optional[List[str]] = Field( + None, description="Required skills for the task" + ) + + +class AgentCommunicationRequest(BaseModel): + message_type: str = Field( + default="notification", + description="Message type (request, response, notification, question)", + ) + content: str = Field(..., description="Message content") + metadata: Dict[str, Any] = Field( + default_factory=dict, description="Additional metadata" + ) + + +class MCPToolRequest(BaseModel): + tool_name: str = Field(..., description="Name of the MCP tool to call") + arguments: Dict[str, Any] = Field( + default_factory=dict, description="Tool arguments" + ) + + +class AgentMCPSetupRequest(BaseModel): + task_id: str = Field(..., description="Task ID for MCP setup") + session_id: Optional[str] = Field(None, description="Optional session ID") + + +class ConversationCreateRequest(BaseModel): + title: str = "New Conversation" + initial_message: Optional[str] = None + context: Optional[Dict[str, Any]] = None + + +class ConversationMessage(BaseModel): + role: str # 'user' or 'agent' + content: str + metadata: Optional[Dict[str, Any]] = None + + +class ChatMessageRequest(BaseModel): + content: str + metadata: Optional[Dict[str, Any]] = None + + +# Model Configuration Models +class ProviderCredentialsRequest(BaseModel): + provider: str = Field( + ..., description="Model provider (anthropic, openai, google, etc.)" + ) + api_key: str = Field(..., description="API key for the provider") + endpoint_url: Optional[str] = Field(None, description="Custom endpoint URL") + additional_config: Dict[str, Any] = Field( + default_factory=dict, description="Additional provider configuration" + ) + + +class AgentModelConfigRequest(BaseModel): + primary_model: str = Field(..., description="Primary model ID") + fallback_models: List[str] = Field( + default_factory=list, description="Fallback model IDs" + ) + temperature: float = Field( + default=0.7, ge=0.0, le=2.0, description="Model temperature" + ) + max_tokens: int = Field( + default=4096, ge=1, le=200000, description="Maximum output tokens" + ) + top_p: float = Field(default=1.0, ge=0.0, le=1.0, description="Top-p sampling") + frequency_penalty: float = Field( + default=0.0, ge=-2.0, le=2.0, description="Frequency penalty" + ) + presence_penalty: float = Field( + default=0.0, ge=-2.0, le=2.0, description="Presence penalty" + ) + custom_instructions: str = Field( + default="", description="Custom instructions for the agent" + ) + use_function_calling: bool = Field( + default=True, description="Enable function calling" + ) + streaming_enabled: bool = Field( + default=True, description="Enable response streaming" + ) + cost_limit_per_task: Optional[float] = Field( + None, ge=0.0, description="Cost limit per task in USD" + ) + + +class TaskCostEstimateRequest(BaseModel): + task_description: str = Field(..., description="Description of the task") + estimated_complexity: str = Field( + default="medium", + description="Estimated complexity (low, medium, high, very_high)", + ) + + +# Response models +class AgentResponse(BaseModel): + agent_id: str + status: str + agent: Dict[str, Any] + + +class TaskResponse(BaseModel): + task_id: str + status: str + + +class CoordinationResponse(BaseModel): + task_id: str + coordination_session_id: Optional[str] = None + status: str + coordination_mode: Optional[str] = None + message: Optional[str] = None + + +# Goals Management API Models +class GoalCreateRequest(BaseModel): + title: str = Field(..., description="Goal title") + description: str = Field(..., description="Goal description") + goal_type: str = Field( + default="business", + description="Goal type (business, technical, growth, operational)", + ) + target_value: Optional[Decimal] = Field( + None, description="Target value (e.g., 100000 for $100K)" + ) + target_unit: Optional[str] = Field( + None, description="Target unit (e.g., 'USD', 'users', '%')" + ) + target_deadline: Optional[date] = Field(None, description="Target completion date") + priority_level: int = Field( + default=5, ge=1, le=10, description="Priority level (1-10)" + ) + success_criteria: Optional[Dict[str, Any]] = Field( + default=None, description="Success criteria" + ) + assigned_teams: Optional[List[str]] = Field( + default=None, description="Assigned team IDs" + ) + goal_owner_agent_id: Optional[str] = Field(None, description="Goal owner agent ID") + stakeholder_agents: Optional[List[str]] = Field( + default=None, description="Stakeholder agent IDs" + ) + tags: Optional[List[str]] = Field(default=None, description="Goal tags") + metadata: Optional[Dict[str, Any]] = Field( + default=None, description="Additional metadata" + ) + + +class GoalUpdateRequest(BaseModel): + progress_percentage: Optional[Decimal] = Field( + None, ge=0, le=100, description="Progress percentage" + ) + current_value: Optional[Decimal] = Field(None, description="Current value") + completion_confidence: Optional[Decimal] = Field( + None, ge=0, le=1, description="Completion confidence" + ) + notes: Optional[str] = Field(None, description="Progress notes") + + +class MilestoneCreateRequest(BaseModel): + title: str = Field(..., description="Milestone title") + description: str = Field(..., description="Milestone description") + target_date: date = Field(..., description="Target completion date") + milestone_type: str = Field(default="deliverable", description="Milestone type") + success_criteria: Optional[Dict[str, Any]] = Field( + default=None, description="Success criteria" + ) + deliverables: Optional[List[Dict[str, Any]]] = Field( + default=None, description="Expected deliverables" + ) + dependencies: Optional[List[Dict[str, Any]]] = Field( + default=None, description="Dependencies" + ) + assigned_teams: Optional[List[str]] = Field( + default=None, description="Assigned teams" + ) + responsible_agent_id: Optional[str] = Field(None, description="Responsible agent") + priority_level: int = Field(default=5, ge=1, le=10, description="Priority level") + weight_in_goal: Optional[Decimal] = Field( + None, ge=0, le=100, description="Weight in goal (%)" + ) + + +class TaskFromMilestoneRequest(BaseModel): + title: str = Field(..., description="Task title") + description: str = Field(..., description="Task description") + task_type: str = Field(default="development", description="Task type") + complexity_level: str = Field(default="medium", description="Complexity level") + estimated_hours: Optional[Decimal] = Field(None, description="Estimated hours") + due_date: Optional[date] = Field(None, description="Due date") + assigned_team_id: Optional[str] = Field(None, description="Assigned team ID") + assigned_agent_id: Optional[str] = Field(None, description="Assigned agent ID") + priority: int = Field(default=5, ge=1, le=10, description="Priority") + requirements: Optional[Dict[str, Any]] = Field( + default=None, description="Requirements" + ) + acceptance_criteria: Optional[List[Dict[str, Any]]] = Field( + default=None, description="Acceptance criteria" + ) + dependencies: Optional[List[Dict[str, Any]]] = Field( + default=None, description="Dependencies" + ) + + +class GoalConversationCreateRequest(BaseModel): + conversation_type: str = Field(default="planning", description="Conversation type") + conversation_title: str = Field(..., description="Conversation title") + initial_context: Optional[Dict[str, Any]] = Field( + default=None, description="Initial context" + ) + participants: Optional[List[Dict[str, Any]]] = Field( + default=None, description="Participants" + ) + + +class ConversationMessageRequest(BaseModel): + message_type: str = Field(default="human", description="Message type") + sender_name: str = Field(..., description="Sender name") + content: str = Field(..., description="Message content") + metadata: Optional[Dict[str, Any]] = Field( + default=None, description="Message metadata" + ) + references: Optional[List[str]] = Field( + default=None, description="Referenced message IDs" + ) + + +class ProgressUpdateRequest(BaseModel): + progress_percentage: Optional[Decimal] = Field( + None, ge=0, le=100, description="Progress percentage" + ) + current_value: Optional[Decimal] = Field(None, description="Current value") + milestone_id: Optional[str] = Field(None, description="Associated milestone ID") + notes: Optional[str] = Field(None, description="Progress notes") + confidence_score: Optional[Decimal] = Field( + None, ge=0, le=1, description="Confidence score" + ) + trigger_alerts: bool = Field(default=True, description="Whether to trigger alerts") + + +@asynccontextmanager +async def lifespan(app: FastAPI): + # Startup + database_url = os.getenv( + "DATABASE_URL", "postgresql://postgres:password@postgres:5432/ai_context" + ) + + app.state.agent_manager = AgentManager(database_url) + app.state.task_queue = TaskQueue() + app.state.context_service = ContextService() + + # Initialize sandbox manager + app.state.sandbox_manager = AgentSandboxManager(database_url) + await app.state.sandbox_manager.start() + + # Start WebSocket manager background cleanup task + await websocket_manager.start() + + # Initialize task execution engine + app.state.task_execution_engine = TaskExecutionEngine(app.state.sandbox_manager) + await app.state.task_execution_engine.start() + + # Initialize multi-agent coordinator + from .multi_agent_coordinator import integrate_multi_agent_coordination + + app.state.multi_agent_coordinator = integrate_multi_agent_coordination( + app.state.task_execution_engine + ) + await app.state.multi_agent_coordinator.start() + + # Initialize knowledge management system + try: + from .context_enhancement_service import ContextEnhancementService + from .knowledge_notification_service import KnowledgeNotificationService + from .knowledge_propagation_engine import KnowledgePropagationEngine + from .organization_rag_manager import OrganizationRAGManager + from .task_knowledge_extractor import TaskKnowledgeExtractor + from .team_knowledge_manager import TeamKnowledgeManager + + app.state.org_rag_manager = OrganizationRAGManager(database_url) + await app.state.org_rag_manager.initialize() + + app.state.team_knowledge_manager = TeamKnowledgeManager(database_url) + await app.state.team_knowledge_manager.initialize() + + app.state.knowledge_propagation_engine = KnowledgePropagationEngine( + database_url, app.state.org_rag_manager, app.state.team_knowledge_manager + ) + await app.state.knowledge_propagation_engine.initialize() + + app.state.notification_service = KnowledgeNotificationService(database_url) + await app.state.notification_service.initialize() + + app.state.task_knowledge_extractor = TaskKnowledgeExtractor( + database_url, + app.state.org_rag_manager, + app.state.team_knowledge_manager, + app.state.knowledge_propagation_engine, + ) + await app.state.task_knowledge_extractor.initialize() + + app.state.context_enhancement_service = ContextEnhancementService( + database_url, app.state.org_rag_manager, app.state.team_knowledge_manager + ) + await app.state.context_enhancement_service.initialize() + + # Initialize knowledge analytics service + from .knowledge_analytics_service import KnowledgeAnalyticsService + + app.state.knowledge_analytics_service = KnowledgeAnalyticsService(database_url) + await app.state.knowledge_analytics_service.initialize() + + logger.info("Knowledge management system initialized successfully") + + except Exception as e: + logger.warning(f"Failed to initialize knowledge management system: {e}") + + # Initialize goals management system + try: + from .goal_conversation_service import GoalConversationService + from .goal_tracking_service import GoalTrackingService + from .goals_management_service import GoalsManagementService + from .milestone_task_engine import MilestoneTaskEngine + + app.state.goals_service = GoalsManagementService(database_url) + await app.state.goals_service.initialize() + + app.state.milestone_task_engine = MilestoneTaskEngine(database_url) + await app.state.milestone_task_engine.initialize() + + app.state.goal_conversation_service = GoalConversationService(database_url) + await app.state.goal_conversation_service.initialize() + + app.state.goal_tracking_service = GoalTrackingService(database_url) + await app.state.goal_tracking_service.initialize() + + logger.info("Goals management system initialized successfully") + + except Exception as e: + logger.warning(f"Failed to initialize goals management system: {e}") + + # Connect components + app.state.task_queue.set_task_execution_engine(app.state.task_execution_engine) + await app.state.agent_manager.set_sandbox_manager(app.state.sandbox_manager) + + # Initialize IzzyAI CEO on startup + try: + await app.state.agent_manager.create_agent( + name="IzzyAI", + role="Digital CEO", + type="executive", + config={ + "model": "claude-sonnet-4-20250514", + "temperature": 0.7, + "tools": [ + "strategic_planning", + "resource_allocation", + "team_management", + ], + }, + ) + except Exception as e: + print(f"Warning: Could not create IzzyAI CEO: {e}") + + yield + + # Shutdown + await app.state.multi_agent_coordinator.stop() + await app.state.task_execution_engine.stop() + await app.state.sandbox_manager.stop() + await app.state.agent_manager.shutdown_all() + await app.state.task_queue.close() + + # Shutdown knowledge management services + try: + if hasattr(app.state, "knowledge_analytics_service"): + await app.state.knowledge_analytics_service.close() + if hasattr(app.state, "context_enhancement_service"): + await app.state.context_enhancement_service.close() + if hasattr(app.state, "task_knowledge_extractor"): + await app.state.task_knowledge_extractor.close() + if hasattr(app.state, "notification_service"): + await app.state.notification_service.close() + if hasattr(app.state, "knowledge_propagation_engine"): + await app.state.knowledge_propagation_engine.close() + if hasattr(app.state, "team_knowledge_manager"): + await app.state.team_knowledge_manager.close() + if hasattr(app.state, "org_rag_manager"): + await app.state.org_rag_manager.close() + logger.info("Knowledge management system shutdown complete") + except Exception as e: + logger.error(f"Error shutting down knowledge management system: {e}") + + # Shutdown goals management services + try: + if hasattr(app.state, "goal_tracking_service"): + await app.state.goal_tracking_service.close() + if hasattr(app.state, "goal_conversation_service"): + await app.state.goal_conversation_service.close() + if hasattr(app.state, "milestone_task_engine"): + await app.state.milestone_task_engine.close() + if hasattr(app.state, "goals_service"): + await app.state.goals_service.close() + logger.info("Goals management system shutdown complete") + except Exception as e: + logger.error(f"Error shutting down goals management system: {e}") + + +app = FastAPI( + title="FuzeAgent Orchestrator API", + description=""" + ## FuzeAgent AI Team Orchestration Platform + + A comprehensive platform for autonomous AI development teams that enables: + + ### 🤖 Autonomous Agent Execution + - **Claude SDK Integration**: Interactive AI development with real-time conversation streaming + - **File Operations Engine**: Safe code changes with human approval workflows + - **Multi-Agent Coordination**: Complex task decomposition and agent collaboration + + ### 🏗️ Core Features + - **Agent Management**: Create, configure, and manage AI development agents + - **Task Orchestration**: Assign and monitor complex development tasks + - **Real-time Monitoring**: WebSocket streaming for live progress updates + - **Human-in-the-Loop**: Seamless approval workflows for critical decisions + + ### 🔗 Integration Capabilities + - **MCP (Model Context Protocol)**: Organizational context for AI agents + - **Git Workflow Management**: Automated repository operations + - **Sandbox Environments**: Isolated development containers + - **Database Integration**: PostgreSQL for persistent storage + + ### 📡 API Categories + - **Agent Management**: Create and manage AI agents + - **Task Execution**: Autonomous task processing + - **File Operations**: Code change management + - **Multi-Agent Coordination**: Team collaboration + - **Real-time Communication**: WebSocket endpoints + - **MCP Integration**: Organizational context tools + - **Goals Management**: Organizational goals, milestones, and task planning + - **Knowledge Management**: RAG system and organizational memory + + **Version**: 2.0.0 (Autonomous Execution) + """, + version="2.0.0", + lifespan=lifespan, + docs_url="/docs", + redoc_url="/redoc", + openapi_tags=[ + {"name": "health", "description": "Health check and system status endpoints"}, + { + "name": "agents", + "description": "AI agent creation, management, and status monitoring", + }, + {"name": "tasks", "description": "Task assignment, execution, and monitoring"}, + { + "name": "autonomous-execution", + "description": "Autonomous task execution with Claude SDK integration", + }, + { + "name": "file-operations", + "description": "File system operations and code change management", + }, + { + "name": "multi-agent-coordination", + "description": "Multi-agent collaboration and task coordination", + }, + { + "name": "real-time", + "description": "WebSocket endpoints for real-time updates", + }, + { + "name": "human-in-loop", + "description": "Human approval workflows and interaction handling", + }, + { + "name": "mcp-integration", + "description": "Model Context Protocol tools and resources", + }, + {"name": "sandboxes", "description": "Sandbox environment management"}, + {"name": "context", "description": "Agent memory and context management"}, + { + "name": "model-configuration", + "description": "AI model configuration and API key management", + }, + { + "name": "knowledge-management", + "description": "Hierarchical knowledge management, RAG, and intelligent notifications", + }, + ], +) + +app.add_middleware( + CORSMiddleware, + allow_origins=[ + "http://localhost:3000", + "http://localhost:3031", + "http://localhost", + ], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +# Include hierarchy router for organizational visualization +app.include_router(hierarchy_router) + + +# Health check endpoint +@app.get( + "/health", + tags=["health"], + summary="Health Check", + description="Check the health status of the FuzeAgent orchestrator service", + response_description="Service health status", +) +async def health_check(): + """ + Health check endpoint that returns the current status of the orchestrator service. + + Returns: + dict: Service health status and basic information + """ + return { + "status": "healthy", + "service": "orchestrator", + "version": "2.0.0", + "features": { + "autonomous_execution": True, + "multi_agent_coordination": True, + "file_operations": True, + "mcp_integration": True, + "real_time_streaming": True, + }, + # Whether GET /openapi.yaml can answer. An image built without its + # contract is DEGRADED, not dead — the probe still passes (no restart + # can conjure a file the image lacks) but the condition is visible to + # anything that looks, instead of surfacing only as a 503 later. + "openapi": "loaded" if _openapi_document() is not None else "unavailable", + } + + +# --------------------------------------------------------------------------- +# The contract, SERVED. +# +# contracts/openapi.yaml describes this orchestrator's real HTTP surface, with +# the curated descriptions and the irreversibility guidance that +# mcp/tools.overrides.yaml narrows. Committing it is not the same as publishing +# it: consumers — the MCP gateway among them — discover the surface over HTTP. +# +# This is NOT /openapi.json. FastAPI generates that from the code at import +# time; it is accurate about shapes and says nothing about which operations +# dispatch an agent that cannot be recalled. Both are served. This one is the +# contract. +# +# The document is read from the IMAGE, never from a mount, so what this endpoint +# publishes is always the contract this build was compiled against. +# --------------------------------------------------------------------------- +_ORCH_DIR = os.path.dirname(os.path.abspath(__file__)) +_OPENAPI_CANDIDATES = [ + p + for p in [ + os.getenv("OPENAPI_SPEC_PATH"), + os.path.join(_ORCH_DIR, "contracts", "openapi.yaml"), + os.path.join(_ORCH_DIR, "..", "..", "contracts", "openapi.yaml"), + ] + if p +] + + +def _openapi_document(): + """Return the OpenAPI document text, or None when the image lacks it.""" + for path in _OPENAPI_CANDIDATES: + try: + with open(path, "r", encoding="utf-8") as fh: + return fh.read() + except OSError: + continue + return None + + +@app.get( + "/openapi.yaml", + tags=["health"], + summary="This OpenAPI Document", + description=( + "Serve contracts/openapi.yaml — the curated contract, as distinct from " + "FastAPI's auto-generated /openapi.json." + ), + include_in_schema=False, +) +async def get_openapi_document(): + doc = _openapi_document() + if doc is None: + logger.error("OpenAPI document not found; tried %s", _OPENAPI_CANDIDATES) + # 503, not 500 and not a crash: the service is otherwise functional and + # no restart can produce a spec the image does not contain. + raise HTTPException( + status_code=503, + detail=( + "openapi_document_unavailable: this image was built without " + "contracts/openapi.yaml. Rebuild with the repo root as the Docker " + "context so the contract is copied in." + ), + ) + return Response(content=doc, media_type="application/yaml") + + +# WebSocket for real-time updates +@app.websocket("/ws") +async def websocket_endpoint(websocket: WebSocket): + await websocket.accept() + try: + while True: + # Send agent updates to UI + updates = await app.state.agent_manager.get_updates() + await websocket.send_json(updates) + await asyncio.sleep(1) + except Exception as e: + print(f"WebSocket error: {e}") + finally: + await websocket.close() + + +# WebSocket for task execution updates +@app.websocket("/ws/tasks/{task_id}") +async def task_websocket_endpoint(websocket: WebSocket, task_id: str): + """WebSocket endpoint for real-time task execution updates""" + await websocket.accept() + try: + while True: + # Get task execution status + try: + status = await app.state.task_queue.get_execution_status(task_id) + await websocket.send_json( + {"type": "status_update", "task_id": task_id, "data": status} + ) + + # If task is completed or failed, send final update and close + if status.get("status") in ["completed", "failed", "cancelled"]: + await websocket.send_json( + { + "type": "task_finished", + "task_id": task_id, + "final_status": status.get("status"), + } + ) + break + + except Exception as e: + await websocket.send_json( + {"type": "error", "task_id": task_id, "error": str(e)} + ) + + await asyncio.sleep(2) # Update every 2 seconds + + except Exception as e: + print(f"Task WebSocket error for {task_id}: {e}") + finally: + await websocket.close() + + +# WebSocket for real-time Claude SDK conversation streaming +@app.websocket("/ws/tasks/{task_id}/conversation") +async def conversation_websocket_endpoint(websocket: WebSocket, task_id: str): + """WebSocket endpoint for real-time Claude SDK conversation streaming""" + await websocket.accept() + try: + # Get execution context + execution = app.state.task_execution_engine.active_executions.get(task_id) + if not execution: + await websocket.send_json( + {"type": "error", "message": f"Task {task_id} not found or not active"} + ) + await websocket.close() + return + + # Wait for Claude SDK session to be available + while not execution.claude_session_id and execution.status not in [ + "completed", + "failed", + "cancelled", + ]: + await asyncio.sleep(1) + + if not execution.claude_session_id: + await websocket.send_json( + { + "type": "error", + "message": "No active Claude SDK session for this task", + } + ) + await websocket.close() + return + + # Stream Claude SDK output + claude_sdk_manager = execution.claude_sdk_manager + if claude_sdk_manager: + try: + async for output_chunk in claude_sdk_manager.stream_output( + execution.claude_session_id + ): + await websocket.send_json( + { + "type": "claude_output", + "task_id": task_id, + "content": output_chunk, + "timestamp": datetime.now().isoformat(), + } + ) + + # Session ended + await websocket.send_json( + { + "type": "conversation_ended", + "task_id": task_id, + "timestamp": datetime.now().isoformat(), + } + ) + + except Exception as e: + await websocket.send_json( + { + "type": "error", + "message": f"Error streaming conversation: {str(e)}", + } + ) + + except Exception as e: + print(f"Conversation WebSocket error for {task_id}: {e}") + finally: + await websocket.close() + + +# WebSocket for file operations streaming +@app.websocket("/ws/tasks/{task_id}/file-operations") +async def file_operations_websocket_endpoint(websocket: WebSocket, task_id: str): + """WebSocket endpoint for real-time file operations updates""" + await websocket.accept() + try: + # Get execution context + execution = app.state.task_execution_engine.active_executions.get(task_id) + if not execution: + await websocket.send_json( + {"type": "error", "message": f"Task {task_id} not found or not active"} + ) + await websocket.close() + return + + file_ops_engine = execution.file_operations_engine + if not file_ops_engine: + await websocket.send_json( + { + "type": "error", + "message": "No file operations engine available for this task", + } + ) + await websocket.close() + return + + last_batch_count = 0 + + while execution.status not in ["completed", "failed", "cancelled"]: + try: + # Get pending operations + pending_operations = file_ops_engine.get_pending_operations() + applied_operations = file_ops_engine.get_applied_operations() + + current_batch_count = len(pending_operations) + len(applied_operations) + + # Send updates if there are new operations + if current_batch_count > last_batch_count: + # Send pending operations + for batch in pending_operations: + # Get diff preview + diffs = await file_ops_engine.get_file_diff_preview( + batch.batch_id + ) + + await websocket.send_json( + { + "type": "pending_operations", + "task_id": task_id, + "batch_id": batch.batch_id, + "description": batch.description, + "requires_approval": batch.requires_approval, + "operations_count": len(batch.operations), + "file_diffs": diffs, + "timestamp": batch.created_at.isoformat(), + } + ) + + # Send applied operations + for batch in applied_operations: + await websocket.send_json( + { + "type": "applied_operations", + "task_id": task_id, + "batch_id": batch.batch_id, + "description": batch.description, + "operations_count": len(batch.operations), + "applied_at": ( + batch.applied_at.isoformat() + if batch.applied_at + else None + ), + "timestamp": batch.created_at.isoformat(), + } + ) + + last_batch_count = current_batch_count + + await asyncio.sleep(1) # Check every second + + except Exception as e: + await websocket.send_json( + { + "type": "error", + "message": f"Error getting file operations: {str(e)}", + } + ) + + # Task completed + await websocket.send_json( + { + "type": "task_completed", + "task_id": task_id, + "final_status": execution.status.value, + "timestamp": datetime.now().isoformat(), + } + ) + + except Exception as e: + print(f"File operations WebSocket error for {task_id}: {e}") + finally: + await websocket.close() + + +# Agent Management Endpoints +@app.post( + "/agents", + tags=["agents"], + summary="Create AI Agent", + description="Create a new AI agent with repository and sandbox settings", + response_model=AgentResponse, +) +async def create_agent(agent_config: AgentCreateRequest, _auth=Depends(require_auth)): + """Create a new AI agent with repository and sandbox settings""" + try: + agent = await app.state.agent_manager.create_agent(**agent_config) + return { + "agent_id": agent.id, + "status": "created", + "agent": { + "id": agent.id, + "name": agent_config.get("name"), + "role": agent_config.get("role"), + "type": agent_config.get("type"), + "repository_settings": agent_config.get("repository_settings", {}), + "sandbox_settings": agent_config.get("sandbox_settings", {}), + "created_at": ( + agent.created_at if hasattr(agent, "created_at") else None + ), + }, + } + except Exception as e: + raise HTTPException(status_code=400, detail=f"Failed to create agent: {str(e)}") + + +@app.get( + "/agents", + tags=["agents"], + summary="List All Agents", + description="Get a list of all AI agents and their current status", +) +async def list_agents(): + """List all agents and their status""" + return await app.state.agent_manager.list_agents() + + +@app.post( + "/agents/{agent_id}/tasks", + tags=["tasks"], + summary="Assign Task to Agent", + description="Assign a specific task to an AI agent", + response_model=TaskResponse, +) +async def assign_task( + agent_id: str = Path(..., description="Agent ID"), + task: TaskCreateRequest = Body(...), + _auth=Depends(require_auth), +): + """Assign a task to an agent""" + task_id = await app.state.task_queue.assign_task(agent_id, task) + return {"task_id": task_id, "status": "assigned"} + + +@app.get("/agents/{agent_id}/status") +async def get_agent_status(agent_id: str): + """Get detailed agent status""" + return await app.state.agent_manager.get_agent_status(agent_id) + + +@app.get("/agents/{agent_id}/tasks") +async def get_agent_tasks(agent_id: str): + """Get tasks assigned to an agent""" + try: + # This would normally query the database for tasks assigned to the agent + # For now, return mock data + return [ + { + "id": "1", + "title": "Strategic Planning Q4 2025", + "description": "Develop comprehensive strategic plan for Q4 2025 expansion", + "status": "completed", + "priority": "high", + "created_at": "2025-08-05T09:00:00Z", + "completed_at": "2025-08-05T17:30:00Z", + }, + { + "id": "2", + "title": "Team Performance Review", + "description": "Conduct quarterly performance review for all team leads", + "status": "in_progress", + "priority": "medium", + "created_at": "2025-08-06T08:00:00Z", + }, + ] + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get agent tasks: {str(e)}" + ) + + +@app.get("/teams") +async def list_teams(): + """List all teams""" + try: + # This would normally query the database for teams + # For now, return mock data + return [ + { + "id": "1", + "name": "Executive Team", + "description": "Strategic leadership and decision making", + }, + { + "id": "2", + "name": "Development Team", + "description": "Frontend, backend, and full-stack development", + }, + { + "id": "3", + "name": "Quality Assurance", + "description": "Testing, quality control, and bug detection", + }, + { + "id": "4", + "name": "DevOps Team", + "description": "Infrastructure, deployment, and system operations", + }, + { + "id": "5", + "name": "Business Team", + "description": "Marketing, sales, and customer relations", + }, + ] + except Exception as e: + raise HTTPException(status_code=500, detail=f"Failed to list teams: {str(e)}") + + +@app.get("/agent-templates") +async def list_agent_templates(): + """List available agent templates""" + try: + return [ + { + "id": "react_developer", + "name": "React Developer", + "description": "Frontend developer specialized in React and TypeScript", + "type": "developer", + "defaultConfig": { + "model": "claude-sonnet-4-20250514", + "temperature": 0.7, + "tools": ["code_generation", "code_review", "debugging", "testing"], + "goal": "Build responsive and performant React applications", + "backstory": "Experienced frontend developer with deep knowledge of React ecosystem", + }, + }, + { + "id": "python_developer", + "name": "Python Developer", + "description": "Backend developer specialized in Python and FastAPI", + "type": "developer", + "defaultConfig": { + "model": "claude-sonnet-4-20250514", + "temperature": 0.7, + "tools": [ + "code_generation", + "api_development", + "database_design", + "testing", + ], + "goal": "Develop robust and scalable backend systems", + "backstory": "Senior Python developer with expertise in FastAPI and databases", + }, + }, + { + "id": "qa_engineer", + "name": "QA Engineer", + "description": "Quality assurance engineer focused on testing and automation", + "type": "qa", + "defaultConfig": { + "model": "claude-sonnet-4-20250514", + "temperature": 0.6, + "tools": [ + "test_automation", + "bug_reporting", + "quality_analysis", + "performance_testing", + ], + "goal": "Ensure high quality and reliability of software products", + "backstory": "Experienced QA engineer with expertise in automated testing frameworks", + }, + }, + { + "id": "devops_engineer", + "name": "DevOps Engineer", + "description": "Infrastructure and deployment specialist", + "type": "devops", + "defaultConfig": { + "model": "claude-sonnet-4-20250514", + "temperature": 0.5, + "tools": [ + "infrastructure_management", + "deployment", + "monitoring", + "security", + ], + "goal": "Maintain reliable and scalable infrastructure", + "backstory": "DevOps engineer with expertise in cloud platforms and CI/CD", + }, + }, + ] + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to list agent templates: {str(e)}" + ) + + +@app.get("/tasks") +async def list_tasks(): + """List all tasks""" + return await app.state.task_queue.list_tasks() + + +@app.get("/tasks/{task_id}") +async def get_task(task_id: str): + """Get task details""" + return await app.state.task_queue.get_task(task_id) + + +# Autonomous Execution Endpoints +@app.post("/agents/from-template") +async def create_agent_from_template(request: dict): + """Create agent from template with repository settings""" + try: + # Extract template data + template_id = request.get("template_id") + name = request.get("name") + team_id = request.get("team_id") + overrides = request.get("overrides", {}) + + # Get template configuration + template_config = await app.state.agent_manager.get_template_config(template_id) + if not template_config: + raise HTTPException( + status_code=404, detail=f"Template {template_id} not found" + ) + + # Build agent configuration + agent_config = { + "name": name, + "role": template_config.get("role", template_id.replace("_", " ").title()), + "type": template_config.get("type", "specialized"), + "template_id": template_id, + "team_id": team_id, + "config": {**template_config.get("config", {}), **overrides}, + "repository_settings": request.get("repository_settings", {}), + "sandbox_settings": { + "base_image": f"fuzeagent/dev-{template_id.split('_')[0]}:latest", + "resource_limits": template_config.get( + "resource_limits", {"memory": "2Gi", "cpu": "1.0", "disk": "10Gi"} + ), + "auto_cleanup": "24h", + }, + } + + # Create agent + agent = await app.state.agent_manager.create_agent(**agent_config) + + return { + "agent_id": agent.id, + "status": "created", + "agent": agent_config, + "template_id": template_id, + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to create agent from template: {str(e)}" + ) + + +@app.get("/templates") +async def get_agent_templates(): + """Get available agent templates""" + return await app.state.agent_manager.get_available_templates() + + +@app.post( + "/tasks/{task_id}/execute", + tags=["autonomous-execution"], + summary="Start Autonomous Task Execution", + description="Begin autonomous execution of a task using Claude SDK integration", + response_model=TaskResponse, +) +async def start_task_execution( + task_id: str = Path(..., description="Task ID to execute") +): + """Start autonomous execution of a task""" + try: + # This will be handled by the TaskExecutionEngine + result = await app.state.task_queue.start_autonomous_execution(task_id) + return {"task_id": task_id, "status": "execution_started", "result": result} + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to start task execution: {str(e)}" + ) + + +@app.get("/tasks/{task_id}/status") +async def get_task_execution_status(task_id: str): + """Get detailed task execution status""" + try: + status = await app.state.task_queue.get_execution_status(task_id) + return status + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get task status: {str(e)}" + ) + + +@app.get("/tasks/{task_id}/iterations") +async def get_task_iterations(task_id: str): + """Get task iteration history""" + try: + iterations = await app.state.task_queue.get_task_iterations(task_id) + return {"task_id": task_id, "iterations": iterations} + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get task iterations: {str(e)}" + ) + + +@app.get("/agents/{agent_id}/sandbox") +async def get_agent_sandbox(agent_id: str): + """Get agent sandbox information""" + try: + sandbox_info = await app.state.agent_manager.get_agent_sandbox(agent_id) + return sandbox_info + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get agent sandbox: {str(e)}" + ) + + +# Additional endpoints for UI support +@app.put("/tasks/{task_id}") +async def update_task(task_id: str, update_data: dict): + """Update task status and result""" + await app.state.task_queue.update_task_status( + task_id=task_id, + status=update_data.get("status"), + result=update_data.get("result"), + ) + return {"status": "updated"} + + +@app.post("/context/interactions") +async def store_interaction(interaction_data: dict): + """Store agent interaction""" + interaction_id = await app.state.context_service.store_interaction( + agent_id=interaction_data.get("agent_id"), + content=interaction_data.get("content"), + metadata=interaction_data.get("metadata", {}), + ) + return {"interaction_id": interaction_id} + + +@app.get("/context") +async def get_context(query: str, agent_id: str = None): + """Get relevant context for a query""" + context = await app.state.context_service.get_context(query, agent_id) + return context + + +@app.get("/agents/{agent_id}/memory") +async def get_agent_memory(agent_id: str, limit: int = 10): + """Get agent memory""" + memory = await app.state.context_service.get_agent_memory(agent_id, limit) + return memory + + +# Agent Conversation Endpoints +@app.get( + "/agents/{agent_id}/conversations", + tags=["conversations"], + summary="Get Agent Conversations", + description="Get all conversations for a specific agent", +) +async def get_agent_conversations(agent_id: str): + """Get all conversations for a specific agent""" + try: + async with get_db_connection() as conn: + conversations = await conn.fetch( + """ + SELECT cs.*, COUNT(ac.id) as message_count, + (SELECT content FROM agent_conversations + WHERE session_id = cs.id + ORDER BY created_at DESC LIMIT 1) as last_message + FROM chat_sessions cs + LEFT JOIN agent_conversations ac ON cs.id = ac.session_id + WHERE cs.agent_id = $1 + GROUP BY cs.id + ORDER BY cs.last_activity DESC + """, + agent_id, + ) + + return [dict(row) for row in conversations] + + except Exception as e: + logger.error(f"Error getting agent conversations: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/agents/{agent_id}/conversations", + tags=["conversations"], + summary="Create New Agent Conversation", + description="Create a new conversation with an agent", +) +async def create_agent_conversation(agent_id: str, request: ConversationCreateRequest): + """Create a new conversation with an agent""" + try: + async with get_db_connection() as conn: + # Create new chat session + session_id = await conn.fetchval( + """ + INSERT INTO chat_sessions (agent_id, session_name, context, status) + VALUES ($1, $2, $3, 'active') + RETURNING id + """, + agent_id, + request.title, + request.context or {}, + ) + + # Add initial message if provided + if request.initial_message: + await conn.execute( + """ + INSERT INTO agent_conversations (session_id, agent_id, message_type, content) + VALUES ($1, $2, 'system', $3) + """, + session_id, + agent_id, + request.initial_message, + ) + + # Get the created conversation + conversation = await conn.fetchrow( + """ + SELECT * FROM chat_sessions WHERE id = $1 + """, + session_id, + ) + + return dict(conversation) + + except Exception as e: + logger.error(f"Error creating agent conversation: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/agents/{agent_id}/conversations/{conversation_id}/messages", + tags=["conversations"], + summary="Get Conversation Messages", + description="Get all messages in a conversation", +) +async def get_conversation_messages(agent_id: str, conversation_id: str): + """Get all messages in a conversation""" + try: + async with get_db_connection() as conn: + messages = await conn.fetch( + """ + SELECT * FROM agent_conversations + WHERE session_id = $1 AND agent_id = $2 + ORDER BY created_at ASC + """, + conversation_id, + agent_id, + ) + + return [dict(row) for row in messages] + + except Exception as e: + logger.error(f"Error getting conversation messages: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/agents/{agent_id}/conversations/{conversation_id}/messages", + tags=["conversations"], + summary="Send Message to Agent", + description="Send a message to an agent in a conversation", +) +async def send_message_to_agent( + agent_id: str, conversation_id: str, request: ChatMessageRequest +): + """Send a message to an agent in a conversation""" + try: + async with get_db_connection() as conn: + # Insert user message + user_message_id = await conn.fetchval( + """ + INSERT INTO agent_conversations (session_id, agent_id, message_type, content, metadata) + VALUES ($1, $2, 'user', $3, $4) + RETURNING id + """, + conversation_id, + agent_id, + request.content, + request.metadata or {}, + ) + + # Update session last activity + await conn.execute( + """ + UPDATE chat_sessions + SET last_activity = CURRENT_TIMESTAMP, message_count = message_count + 1 + WHERE id = $1 + """, + conversation_id, + ) + + # TODO: Here we would integrate with the actual agent to generate a response + # For now, return a simple acknowledgment + + return { + "id": str(user_message_id), + "status": "sent", + "message": "Message sent to agent", + } + + except Exception as e: + logger.error(f"Error sending message to agent: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get("/knowledge/search") +async def search_knowledge(query: str, limit: int = 10): + """Search knowledge base""" + results = await app.state.context_service.search_knowledge(query, limit) + return results + + +# Human-in-the-loop endpoints +@app.post( + "/tasks/{task_id}/human-response", + tags=["human-in-loop"], + summary="Submit Human Response", + description="Submit human response to a task question or approval request", +) +async def submit_human_response( + task_id: str = Path(..., description="Task ID"), + response_data: HumanResponseRequest = Body(...), +): + """Submit human response to a task question""" + try: + response = response_data.get("response", "") + if not response: + raise HTTPException(status_code=400, detail="Response cannot be empty") + + success = await app.state.task_queue.handle_human_response(task_id, response) + + if success: + return {"status": "success", "message": "Human response submitted"} + else: + raise HTTPException( + status_code=404, + detail="Task not found or not waiting for human response", + ) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to submit human response: {str(e)}" + ) + + +@app.post("/tasks/{task_id}/cancel") +async def cancel_task_execution(task_id: str): + """Cancel autonomous execution of a task""" + try: + success = await app.state.task_queue.cancel_task_execution(task_id) + + if success: + return {"status": "cancelled", "message": "Task execution cancelled"} + else: + raise HTTPException(status_code=404, detail="Task not found or not running") + + except Exception as e: + raise HTTPException(status_code=500, detail=f"Failed to cancel task: {str(e)}") + + +@app.get("/tasks/{task_id}/messages") +async def get_task_messages(task_id: str): + """Get task messages and chat history""" + try: + # This would integrate with the HumanInTheLoopHandler when implemented + # For now, return iteration history which includes human interactions + iterations = await app.state.task_queue.get_task_iterations(task_id) + + messages = [] + for iteration in iterations: + if iteration.get("human_question"): + messages.append( + { + "type": "agent_question", + "content": iteration["human_question"], + "timestamp": iteration["started_at"], + "iteration": iteration["iteration_number"], + } + ) + + if iteration.get("human_response"): + messages.append( + { + "type": "human_response", + "content": iteration["human_response"], + "timestamp": iteration["completed_at"] + or iteration["started_at"], + "iteration": iteration["iteration_number"], + } + ) + + return {"task_id": task_id, "messages": messages} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get task messages: {str(e)}" + ) + + +# Sandbox management endpoints +@app.get("/sandboxes") +async def list_sandboxes(agent_id: str = None, status: str = None): + """List active sandboxes""" + try: + from .sandbox_manager import SandboxStatus + + sandbox_status = None + if status: + try: + sandbox_status = SandboxStatus(status) + except ValueError: + raise HTTPException(status_code=400, detail=f"Invalid status: {status}") + + sandboxes = await app.state.sandbox_manager.list_sandboxes( + agent_id=agent_id, status=sandbox_status + ) + + return { + "sandboxes": [ + { + "sandbox_id": s.sandbox_id, + "agent_id": s.agent_id, + "task_id": s.task_id, + "status": s.status.value, + "workspace_path": s.workspace_path, + "created_at": s.created_at.isoformat(), + "resource_limits": s.resource_limits, + } + for s in sandboxes + ] + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to list sandboxes: {str(e)}" + ) + + +@app.post("/sandboxes/{sandbox_id}/execute") +async def execute_command_in_sandbox(sandbox_id: str, command_data: dict): + """Execute a command in a sandbox""" + try: + command = command_data.get("command") + working_dir = command_data.get("working_dir") + + if not command: + raise HTTPException(status_code=400, detail="Command is required") + + result = await app.state.sandbox_manager.execute_command( + sandbox_id=sandbox_id, command=command, working_dir=working_dir + ) + + return result + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to execute command: {str(e)}" + ) + + +@app.delete("/sandboxes/{sandbox_id}") +async def destroy_sandbox(sandbox_id: str): + """Destroy a sandbox""" + try: + await app.state.sandbox_manager.destroy_sandbox(sandbox_id) + return {"status": "destroyed", "sandbox_id": sandbox_id} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to destroy sandbox: {str(e)}" + ) + + +# Agent registration and communication endpoints +@app.post("/agents/{agent_id}/register") +async def register_agent(agent_id: str, registration_data: dict): + """Register an agent running in a sandbox container""" + try: + # Store agent registration info + # This would typically update the agent's status and capabilities + return { + "status": "registered", + "agent_id": agent_id, + "registered_at": datetime.now().isoformat(), + } + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to register agent: {str(e)}" + ) + + +@app.get("/agents/{agent_id}/next-task") +async def get_next_task_for_agent(agent_id: str): + """Get the next task for an agent to execute""" + try: + # Find pending tasks assigned to this agent + tasks = await app.state.task_queue.get_agent_tasks(agent_id) + pending_tasks = [t for t in tasks if t.get("status") == "pending"] + + if pending_tasks: + # Return the first pending task + task = pending_tasks[0] + # Update status to 'assigned' to prevent double assignment + await app.state.task_queue.update_task_status(task["id"], "assigned") + return task + else: + # No tasks available + return None, 204 + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get next task: {str(e)}" + ) + + +@app.post("/agents/{agent_id}/error") +async def report_agent_error(agent_id: str, error_data: dict): + """Report an error from an agent""" + try: + # Log the error and update agent status + logger.error(f"Agent {agent_id} reported error: {error_data.get('error')}") + + # You might want to store this in a database or alerting system + return {"status": "error_logged", "agent_id": agent_id} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to log agent error: {str(e)}" + ) + + +# Conversation management endpoints +@app.get("/tasks/{task_id}/conversation") +async def get_task_conversation(task_id: str, iteration: int = None, limit: int = 100): + """Get conversation history for a task""" + try: + conversation_history = await app.state.task_execution_engine.conversation_manager.get_conversation_history( + task_id=task_id, iteration_number=iteration, limit=limit + ) + + return {"task_id": task_id, "conversation": conversation_history} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get conversation: {str(e)}" + ) + + +@app.get("/tasks/{task_id}/conversation/summary") +async def get_conversation_summary(task_id: str): + """Get conversation summary with statistics""" + try: + summary = await app.state.task_execution_engine.conversation_manager.get_conversation_summary( + task_id + ) + return summary + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get conversation summary: {str(e)}" + ) + + +@app.get("/tasks/{task_id}/code-generations") +async def get_task_code_generations( + task_id: str, iteration: int = None, file_type: str = None +): + """Get code generations for a task""" + try: + code_generations = await app.state.task_execution_engine.conversation_manager.get_code_generations( + task_id=task_id, iteration_number=iteration, file_type=file_type + ) + + return {"task_id": task_id, "code_generations": code_generations} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get code generations: {str(e)}" + ) + + +@app.get("/agents/{agent_id}/performance") +async def get_agent_performance(agent_id: str, hours: int = 24): + """Get agent performance metrics""" + try: + metrics = await app.state.task_execution_engine.conversation_manager.get_agent_performance_metrics( + agent_id=agent_id, time_range_hours=hours + ) + + return {"agent_id": agent_id, "time_range_hours": hours, "metrics": metrics} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get agent performance: {str(e)}" + ) + + +# File Operations Endpoints +@app.get( + "/tasks/{task_id}/file-operations", + tags=["file-operations"], + summary="Get File Operations", + description="Get file operations for a task with optional status filtering", +) +async def get_task_file_operations( + task_id: str = Path(..., description="Task ID"), + status: Optional[str] = Query( + None, description="Filter by status (pending, applied)" + ), +): + """Get file operations for a task""" + try: + execution = app.state.task_execution_engine.active_executions.get(task_id) + if not execution or not execution.file_operations_engine: + raise HTTPException( + status_code=404, detail="Task not found or no file operations available" + ) + + file_ops_engine = execution.file_operations_engine + + if status == "pending": + operations = file_ops_engine.get_pending_operations() + elif status == "applied": + operations = file_ops_engine.get_applied_operations() + else: + # Get all operations + pending = file_ops_engine.get_pending_operations() + applied = file_ops_engine.get_applied_operations() + operations = pending + applied + + # Convert to dict format + operations_data = [] + for batch in operations: + operations_data.append( + { + "batch_id": batch.batch_id, + "task_id": batch.task_id, + "agent_id": batch.agent_id, + "description": batch.description, + "requires_approval": batch.requires_approval, + "approval_status": batch.approval_status.value, + "operations_count": len(batch.operations), + "created_at": batch.created_at.isoformat(), + "applied_at": ( + batch.applied_at.isoformat() if batch.applied_at else None + ), + } + ) + + return {"task_id": task_id, "operations": operations_data} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get file operations: {str(e)}" + ) + + +@app.get("/tasks/{task_id}/file-operations/{batch_id}/preview") +async def get_file_operations_preview(task_id: str, batch_id: str): + """Get preview of file changes for a batch""" + try: + execution = app.state.task_execution_engine.active_executions.get(task_id) + if not execution or not execution.file_operations_engine: + raise HTTPException( + status_code=404, detail="Task not found or no file operations available" + ) + + file_ops_engine = execution.file_operations_engine + diffs = await file_ops_engine.get_file_diff_preview(batch_id) + + return {"task_id": task_id, "batch_id": batch_id, "file_diffs": diffs} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get file preview: {str(e)}" + ) + + +@app.post( + "/tasks/{task_id}/file-operations/{batch_id}/approve", + tags=["file-operations", "human-in-loop"], + summary="Approve File Operations", + description="Approve or reject file operations from Claude SDK", +) +async def approve_file_operations( + task_id: str = Path(..., description="Task ID"), + batch_id: str = Path(..., description="Batch ID"), + approval_data: FileOperationApprovalRequest = Body(...), +): + """Approve or reject file operations""" + try: + approved = approval_data.get("approved", False) + + execution = app.state.task_execution_engine.active_executions.get(task_id) + if not execution or not execution.file_operations_engine: + raise HTTPException( + status_code=404, detail="Task not found or no file operations available" + ) + + file_ops_engine = execution.file_operations_engine + success = await file_ops_engine.approve_operations(batch_id, approved) + + if success: + # Also notify Claude SDK if there's an active session + if execution.claude_sdk_manager and execution.claude_session_id: + await execution.claude_sdk_manager.approve_file_operations( + execution.claude_session_id, batch_id, approved + ) + + return { + "task_id": task_id, + "batch_id": batch_id, + "approved": approved, + "status": "success", + } + else: + raise HTTPException(status_code=400, detail="Failed to process approval") + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to approve file operations: {str(e)}" + ) + + +@app.post("/tasks/{task_id}/file-operations/{batch_id}/rollback") +async def rollback_file_operations(task_id: str, batch_id: str): + """Rollback applied file operations""" + try: + execution = app.state.task_execution_engine.active_executions.get(task_id) + if not execution or not execution.file_operations_engine: + raise HTTPException( + status_code=404, detail="Task not found or no file operations available" + ) + + file_ops_engine = execution.file_operations_engine + success = await file_ops_engine.rollback_operations(batch_id) + + if success: + return {"task_id": task_id, "batch_id": batch_id, "status": "rolled_back"} + else: + raise HTTPException(status_code=400, detail="Failed to rollback operations") + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to rollback file operations: {str(e)}" + ) + + +# Claude SDK Session Management Endpoints +@app.get("/tasks/{task_id}/claude-session") +async def get_claude_session_status(task_id: str): + """Get Claude SDK session status for a task""" + try: + execution = app.state.task_execution_engine.active_executions.get(task_id) + if ( + not execution + or not execution.claude_sdk_manager + or not execution.claude_session_id + ): + raise HTTPException( + status_code=404, detail="No active Claude SDK session for this task" + ) + + status = await execution.claude_sdk_manager.get_session_status( + execution.claude_session_id + ) + return status + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get Claude session status: {str(e)}" + ) + + +@app.post("/tasks/{task_id}/claude-session/input") +async def send_claude_session_input(task_id: str, input_data: dict): + """Send input to Claude SDK session""" + try: + user_input = input_data.get("input", "") + if not user_input: + raise HTTPException(status_code=400, detail="Input cannot be empty") + + execution = app.state.task_execution_engine.active_executions.get(task_id) + if ( + not execution + or not execution.claude_sdk_manager + or not execution.claude_session_id + ): + raise HTTPException( + status_code=404, detail="No active Claude SDK session for this task" + ) + + success = await execution.claude_sdk_manager.send_input( + execution.claude_session_id, user_input + ) + + if success: + return {"task_id": task_id, "status": "input_sent", "input": user_input} + else: + raise HTTPException( + status_code=400, detail="Failed to send input to Claude session" + ) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to send Claude session input: {str(e)}" + ) + + +# MCP Integration Endpoints +@app.get("/mcp/tools") +async def get_mcp_tools(): + """Get available MCP tools""" + try: + from .mcp_integration import FuzeAgentMCPServer + + mcp_server = FuzeAgentMCPServer() + tools = [ + { + "name": tool.name, + "description": tool.description, + "input_schema": tool.input_schema, + } + for tool in mcp_server.tools + ] + + return {"tools": tools} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get MCP tools: {str(e)}" + ) + + +@app.post( + "/mcp/call-tool", + tags=["mcp-integration"], + summary="Call MCP Tool", + description="Execute an MCP tool to access organizational context", +) +async def call_mcp_tool(tool_request: MCPToolRequest = Body(...)): + """Call an MCP tool""" + try: + from .mcp_integration import FuzeAgentMCPServer + + tool_name = tool_request.get("tool_name") + arguments = tool_request.get("arguments", {}) + + if not tool_name: + raise HTTPException(status_code=400, detail="tool_name is required") + + mcp_server = FuzeAgentMCPServer() + result = await mcp_server.handle_tool_call(tool_name, arguments) + + return result + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to call MCP tool: {str(e)}" + ) + + +@app.get("/mcp/resources") +async def get_mcp_resources(): + """Get available MCP resources""" + try: + from .mcp_integration import FuzeAgentMCPServer + + mcp_server = FuzeAgentMCPServer() + resources = [ + { + "uri": resource.uri, + "name": resource.name, + "description": resource.description, + "mime_type": resource.mime_type, + } + for resource in mcp_server.resources + ] + + return {"resources": resources} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get MCP resources: {str(e)}" + ) + + +@app.get("/mcp/resource") +async def get_mcp_resource(uri: str): + """Get an MCP resource by URI""" + try: + from .mcp_integration import FuzeAgentMCPServer + + if not uri: + raise HTTPException(status_code=400, detail="uri parameter is required") + + mcp_server = FuzeAgentMCPServer() + resource = await mcp_server.handle_resource_request(uri) + + return resource + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get MCP resource: {str(e)}" + ) + + +@app.get("/tasks/{task_id}/mcp-context") +async def get_task_mcp_context(task_id: str): + """Get MCP context for a task""" + try: + from .mcp_integration import FuzeAgentMCPServer, MCPClaudeIntegration + + execution = app.state.task_execution_engine.active_executions.get(task_id) + if not execution: + raise HTTPException(status_code=404, detail="Task not found or not active") + + mcp_server = FuzeAgentMCPServer() + mcp_integration = MCPClaudeIntegration(mcp_server) + + session_id = execution.claude_session_id or f"session-{task_id}" + context = await mcp_integration.get_session_context( + session_id=session_id, agent_id=execution.agent_id, task_id=task_id + ) + + return context + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get MCP context: {str(e)}" + ) + + +@app.post( + "/agents/{agent_id}/mcp-setup", + tags=["mcp-integration"], + summary="Setup Agent MCP Integration", + description="Configure MCP integration for an AI agent", +) +async def setup_agent_mcp( + agent_id: str = Path(..., description="Agent ID"), + setup_data: AgentMCPSetupRequest = Body(...), +): + """Set up MCP integration for an agent""" + try: + from .mcp_integration import FuzeAgentMCPServer, MCPClaudeIntegration + + task_id = setup_data.get("task_id") + session_id = setup_data.get("session_id") + + if not task_id: + raise HTTPException(status_code=400, detail="task_id is required") + + mcp_server = FuzeAgentMCPServer() + mcp_integration = MCPClaudeIntegration(mcp_server) + + # Set up MCP for Claude session + mcp_config = await mcp_integration.setup_claude_session_mcp( + session_id=session_id or f"session-{task_id}", + agent_id=agent_id, + task_id=task_id, + ) + + return { + "agent_id": agent_id, + "task_id": task_id, + "mcp_config": mcp_config, + "status": "mcp_configured", + } + + except Exception as e: + raise HTTPException(status_code=500, detail=f"Failed to setup MCP: {str(e)}") + + +# Multi-Agent Coordination Endpoints +@app.post( + "/tasks/{task_id}/coordinate", + tags=["multi-agent-coordination"], + summary="Initiate Multi-Agent Coordination", + description="Initiate multi-agent coordination for complex tasks", + response_model=CoordinationResponse, +) +async def initiate_task_coordination( + task_id: str = Path(..., description="Task ID to coordinate"), + coordination_request: CoordinationRequest = Body(...), +): + """Initiate multi-agent coordination for a complex task""" + try: + from .multi_agent_coordinator import CoordinationMode + + coordination_mode = coordination_request.get( + "coordination_mode", "collaborative" + ) + required_agents = coordination_request.get("required_agents") + required_skills = coordination_request.get("required_skills") + + # Validate coordination mode + try: + coord_mode = CoordinationMode(coordination_mode) + except ValueError: + raise HTTPException( + status_code=400, + detail=f"Invalid coordination mode: {coordination_mode}", + ) + + # Get multi-agent coordinator + coordinator = getattr( + app.state.task_execution_engine, "multi_agent_coordinator", None + ) + if not coordinator: + raise HTTPException( + status_code=503, detail="Multi-agent coordination not available" + ) + + # Initiate coordination + session_id = await coordinator.initiate_coordination( + task_id=task_id, + coordination_mode=coord_mode, + required_agents=required_agents, + required_skills=required_skills, + ) + + if session_id: + return { + "task_id": task_id, + "coordination_session_id": session_id, + "status": "coordination_initiated", + "coordination_mode": coordination_mode, + } + else: + return { + "task_id": task_id, + "status": "coordination_not_needed", + "message": "Task does not require multi-agent coordination", + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to initiate coordination: {str(e)}" + ) + + +@app.get("/coordination/{session_id}") +async def get_coordination_status(session_id: str): + """Get status of a coordination session""" + try: + coordinator = getattr( + app.state.task_execution_engine, "multi_agent_coordinator", None + ) + if not coordinator: + raise HTTPException( + status_code=503, detail="Multi-agent coordination not available" + ) + + status = await coordinator.get_coordination_status(session_id) + + if status: + return status + else: + raise HTTPException( + status_code=404, detail="Coordination session not found" + ) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get coordination status: {str(e)}" + ) + + +@app.post("/coordination/{session_id}/cancel") +async def cancel_coordination(session_id: str): + """Cancel a coordination session""" + try: + coordinator = getattr( + app.state.task_execution_engine, "multi_agent_coordinator", None + ) + if not coordinator: + raise HTTPException( + status_code=503, detail="Multi-agent coordination not available" + ) + + success = await coordinator.cancel_coordination(session_id) + + if success: + return {"coordination_session_id": session_id, "status": "cancelled"} + else: + raise HTTPException( + status_code=404, detail="Coordination session not found" + ) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to cancel coordination: {str(e)}" + ) + + +@app.post("/agents/{from_agent_id}/communicate/{to_agent_id}") +async def send_agent_communication( + from_agent_id: str, to_agent_id: str, communication_data: dict +): + """Send communication between agents""" + try: + message_type = communication_data.get("message_type", "notification") + content = communication_data.get("content", "") + metadata = communication_data.get("metadata", {}) + + if not content: + raise HTTPException(status_code=400, detail="Content cannot be empty") + + coordinator = getattr( + app.state.task_execution_engine, "multi_agent_coordinator", None + ) + if not coordinator: + raise HTTPException( + status_code=503, detail="Multi-agent coordination not available" + ) + + communication_id = await coordinator.send_agent_communication( + from_agent_id=from_agent_id, + to_agent_id=to_agent_id, + message_type=message_type, + content=content, + metadata=metadata, + ) + + return { + "communication_id": communication_id, + "from_agent_id": from_agent_id, + "to_agent_id": to_agent_id, + "status": "sent", + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to send agent communication: {str(e)}" + ) + + +@app.get("/coordination/active") +async def get_active_coordinations(): + """Get all active coordination sessions""" + try: + coordinator = getattr( + app.state.task_execution_engine, "multi_agent_coordinator", None + ) + if not coordinator: + raise HTTPException( + status_code=503, detail="Multi-agent coordination not available" + ) + + active_sessions = [] + for session_id in coordinator.active_sessions.keys(): + status = await coordinator.get_coordination_status(session_id) + if status: + active_sessions.append(status) + + return {"active_coordinations": active_sessions, "count": len(active_sessions)} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get active coordinations: {str(e)}" + ) + + +# WebSocket for coordination updates +@app.websocket("/ws/coordination/{session_id}") +async def coordination_websocket_endpoint(websocket: WebSocket, session_id: str): + """WebSocket endpoint for real-time coordination updates""" + await websocket.accept() + try: + coordinator = getattr( + app.state.task_execution_engine, "multi_agent_coordinator", None + ) + if not coordinator: + await websocket.send_json( + {"type": "error", "message": "Multi-agent coordination not available"} + ) + await websocket.close() + return + + # Monitor coordination session + while True: + try: + status = await coordinator.get_coordination_status(session_id) + if status: + await websocket.send_json( + { + "type": "coordination_update", + "session_id": session_id, + "data": status, + "timestamp": datetime.now().isoformat(), + } + ) + + # If coordination is completed or failed, send final update + if status.get("status") in ["completed", "failed", "cancelled"]: + await websocket.send_json( + { + "type": "coordination_finished", + "session_id": session_id, + "final_status": status.get("status"), + "timestamp": datetime.now().isoformat(), + } + ) + break + else: + await websocket.send_json( + { + "type": "error", + "message": f"Coordination session {session_id} not found", + } + ) + break + + await asyncio.sleep(3) # Update every 3 seconds + + except Exception as e: + await websocket.send_json( + { + "type": "error", + "message": f"Error monitoring coordination: {str(e)}", + } + ) + + except Exception as e: + print(f"Coordination WebSocket error for {session_id}: {e}") + finally: + await websocket.close() + + +# --------------------------------------------------------------------------- +# Agent relay WebSocket (Track 4) +# --------------------------------------------------------------------------- +@app.websocket("/agent-relay/{agent_id}") +async def agent_relay_endpoint(websocket: WebSocket, agent_id: str): + """ + Agent pods connect here to stream their session output. + Dashboard clients connect here to watch a specific agent's session. + Both use the same endpoint — first JSON message determines role: + {"role": "agent"} -> agent pod streaming output + {"role": "subscriber"} -> human dashboard watcher (default) + """ + await websocket.accept() + role = None + try: + init_msg = await websocket.receive_json() + role = init_msg.get("role", "subscriber") + + if role == "agent": + # Stream from agent pod to all subscribers + async for data in websocket.iter_json(): + msg = {"agentId": agent_id, **data} + dead = [] + for sub in list(agent_relay_subscribers[agent_id]): + try: + await sub.send_json(msg) + except Exception: + dead.append(sub) + for d in dead: + agent_relay_subscribers[agent_id].remove(d) + else: + # Human dashboard subscriber — wait for messages from agent + agent_relay_subscribers[agent_id].append(websocket) + await websocket.receive_text() # keep alive until disconnect + except WebSocketDisconnect: + pass + except Exception as e: + logger.warning(f"agent-relay {agent_id}: {e}") + finally: + subs = agent_relay_subscribers.get(agent_id, []) + if role != "agent" and websocket in subs: + subs.remove(websocket) + + +# Model Configuration and API Key Management Endpoints +@app.post( + "/organizations/{organization_id}/providers/{provider}/credentials", + tags=["model-configuration"], + summary="Store Provider API Credentials", + description="Store encrypted API credentials for a model provider", +) +async def store_provider_credentials( + organization_id: str = Path(..., description="Organization ID"), + provider: str = Path(..., description="Provider name"), + credentials: ProviderCredentialsRequest = Body(...), +): + """Store encrypted API credentials for a model provider at organization level""" + try: + from .model_configuration import ModelProvider, model_config_manager + + # Validate provider + try: + provider_enum = ModelProvider(provider) + except ValueError: + raise HTTPException( + status_code=400, detail=f"Unsupported provider: {provider}" + ) + + success = await model_config_manager.store_provider_credentials( + organization_id=organization_id, + provider=provider_enum, + api_key=credentials.api_key, + endpoint_url=credentials.endpoint_url, + additional_config=credentials.additional_config, + ) + + if success: + return { + "organization_id": organization_id, + "provider": provider, + "status": "credentials_stored", + "message": "API credentials stored successfully", + } + else: + raise HTTPException(status_code=500, detail="Failed to store credentials") + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to store provider credentials: {str(e)}" + ) + + +@app.get( + "/organizations/{organization_id}/models", + tags=["model-configuration"], + summary="Get Available Models", + description="Get available AI models for an organization", +) +async def get_available_models( + organization_id: str = Path(..., description="Organization ID"), + provider: Optional[str] = Query(None, description="Filter by provider"), + capabilities: Optional[str] = Query( + None, description="Filter by capabilities (comma-separated)" + ), +): + """Get available AI models with provider credential validation""" + try: + from .model_configuration import ( + ModelCapability, + ModelProvider, + model_config_manager, + ) + + provider_filter = None + if provider: + try: + provider_filter = ModelProvider(provider) + except ValueError: + raise HTTPException( + status_code=400, detail=f"Invalid provider: {provider}" + ) + + capabilities_filter = None + if capabilities: + try: + capabilities_filter = [ + ModelCapability(cap.strip()) for cap in capabilities.split(",") + ] + except ValueError as e: + raise HTTPException( + status_code=400, detail=f"Invalid capability: {str(e)}" + ) + + models = await model_config_manager.get_available_models( + organization_id=organization_id, + provider=provider_filter, + capabilities=capabilities_filter, + ) + + return { + "organization_id": organization_id, + "models": models, + "count": len(models), + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get available models: {str(e)}" + ) + + +@app.post( + "/agents/{agent_id}/model-configuration", + tags=["model-configuration"], + summary="Configure Agent Model Settings", + description="Configure model settings and preferences for an agent", +) +async def configure_agent_model( + agent_id: str = Path(..., description="Agent ID"), + config: AgentModelConfigRequest = Body(...), +): + """Configure model settings for an AI agent""" + try: + from .model_configuration import AgentModelConfig, model_config_manager + + agent_config = AgentModelConfig( + agent_id=agent_id, + primary_model=config.primary_model, + fallback_models=config.fallback_models, + temperature=config.temperature, + max_tokens=config.max_tokens, + top_p=config.top_p, + frequency_penalty=config.frequency_penalty, + presence_penalty=config.presence_penalty, + custom_instructions=config.custom_instructions, + use_function_calling=config.use_function_calling, + streaming_enabled=config.streaming_enabled, + cost_limit_per_task=config.cost_limit_per_task, + ) + + success = await model_config_manager.configure_agent_model( + agent_id, agent_config + ) + + if success: + return { + "agent_id": agent_id, + "status": "configured", + "primary_model": config.primary_model, + "fallback_models": config.fallback_models, + } + else: + raise HTTPException( + status_code=500, detail="Failed to configure agent model" + ) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to configure agent model: {str(e)}" + ) + + +@app.get( + "/agents/{agent_id}/model-configuration", + tags=["model-configuration"], + summary="Get Agent Model Configuration", + description="Get current model configuration for an agent", +) +async def get_agent_model_configuration( + agent_id: str = Path(..., description="Agent ID") +): + """Get model configuration for an AI agent""" + try: + from .model_configuration import model_config_manager + + config = await model_config_manager.get_agent_model_config(agent_id) + + if config: + return { + "agent_id": agent_id, + "configuration": { + "primary_model": config.primary_model, + "fallback_models": config.fallback_models, + "temperature": config.temperature, + "max_tokens": config.max_tokens, + "top_p": config.top_p, + "frequency_penalty": config.frequency_penalty, + "presence_penalty": config.presence_penalty, + "custom_instructions": config.custom_instructions, + "use_function_calling": config.use_function_calling, + "streaming_enabled": config.streaming_enabled, + "cost_limit_per_task": config.cost_limit_per_task, + "created_at": config.created_at.isoformat(), + "updated_at": config.updated_at.isoformat(), + }, + } + else: + raise HTTPException( + status_code=404, detail="Agent model configuration not found" + ) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get agent model configuration: {str(e)}" + ) + + +@app.post( + "/agents/{agent_id}/tasks/cost-estimate", + tags=["model-configuration"], + summary="Estimate Task Cost", + description="Estimate the cost of executing a task with the agent's model configuration", +) +async def estimate_task_cost( + agent_id: str = Path(..., description="Agent ID"), + request: TaskCostEstimateRequest = Body(...), +): + """Estimate cost for task execution based on agent's model configuration""" + try: + from .model_configuration import model_config_manager + + estimate = await model_config_manager.estimate_task_cost( + agent_id=agent_id, + task_description=request.task_description, + estimated_complexity=request.estimated_complexity, + ) + + return estimate + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to estimate task cost: {str(e)}" + ) + + +@app.get( + "/organizations/{organization_id}/model-usage", + tags=["model-configuration"], + summary="Get Model Usage Statistics", + description="Get model usage statistics and costs for an organization", +) +async def get_organization_model_usage( + organization_id: str = Path(..., description="Organization ID"), + days: int = Query(30, ge=1, le=365, description="Number of days to analyze"), +): + """Get model usage statistics and costs for an organization""" + try: + from .model_configuration import model_config_manager + + usage = await model_config_manager.get_organization_model_usage( + organization_id=organization_id, days=days + ) + + return usage + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get model usage: {str(e)}" + ) + + +@app.get( + "/agents/{agent_id}/model-recommendations", + tags=["model-configuration"], + summary="Get Model Recommendations", + description="Get model recommendations for an agent based on task capabilities", +) +async def get_model_recommendations( + agent_id: str = Path(..., description="Agent ID"), + capabilities: str = Query( + ..., description="Required capabilities (comma-separated)" + ), + cost_limit: Optional[float] = Query( + None, ge=0.0, description="Maximum cost limit in USD" + ), +): + """Get model recommendations based on task capabilities and cost constraints""" + try: + from .model_configuration import ModelCapability, model_config_manager + + # Parse capabilities + try: + capability_list = [ + ModelCapability(cap.strip()) for cap in capabilities.split(",") + ] + except ValueError as e: + raise HTTPException(status_code=400, detail=f"Invalid capability: {str(e)}") + + recommended_model = await model_config_manager.get_model_for_task( + agent_id=agent_id, task_capabilities=capability_list, cost_limit=cost_limit + ) + + if recommended_model: + return { + "agent_id": agent_id, + "recommended_model": recommended_model, + "capabilities": capabilities, + "cost_limit": cost_limit, + } + else: + return { + "agent_id": agent_id, + "recommended_model": None, + "message": "No suitable model found for the specified requirements", + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get model recommendations: {str(e)}" + ) + + +# Knowledge Management and Notification Endpoints + + +@app.get( + "/knowledge/notifications/{recipient_type}/{recipient_id}", + tags=["knowledge-management"], + summary="Get Knowledge Notifications", + description="Get notifications about knowledge updates, conflicts, and opportunities", +) +async def get_knowledge_notifications( + recipient_type: str = Path( + ..., description="Recipient type (agent, team, organization)" + ), + recipient_id: str = Path(..., description="Recipient ID"), + limit: int = Query(20, ge=1, le=100, description="Maximum notifications to return"), + status_filter: Optional[str] = Query( + None, description="Filter by status (unread, read, acknowledged)" + ), + notification_type_filter: Optional[str] = Query( + None, description="Filter by type (comma-separated)" + ), +): + """Get knowledge notifications for a recipient""" + try: + from .knowledge_notification_service import ( + KnowledgeNotificationService, + NotificationStatus, + NotificationType, + ) + + # Initialize notification service if not already done + if not hasattr(app.state, "notification_service"): + database_url = os.getenv( + "DATABASE_URL", + "postgresql://postgres:password@postgres:5432/ai_context", + ) + app.state.notification_service = KnowledgeNotificationService(database_url) + await app.state.notification_service.initialize() + + # Parse filters + status_filters = None + if status_filter: + try: + status_filters = [ + NotificationStatus(s.strip()) for s in status_filter.split(",") + ] + except ValueError as e: + raise HTTPException( + status_code=400, detail=f"Invalid status filter: {str(e)}" + ) + + type_filters = None + if notification_type_filter: + try: + type_filters = [ + NotificationType(t.strip()) + for t in notification_type_filter.split(",") + ] + except ValueError as e: + raise HTTPException( + status_code=400, + detail=f"Invalid notification type filter: {str(e)}", + ) + + notifications = ( + await app.state.notification_service.get_notifications_for_recipient( + recipient_type=recipient_type, + recipient_id=recipient_id, + limit=limit, + status_filter=status_filters, + notification_type_filter=type_filters, + ) + ) + + return { + "recipient_type": recipient_type, + "recipient_id": recipient_id, + "notifications": [ + { + "id": n.id, + "notification_type": n.notification_type.value, + "title": n.title, + "message": n.message, + "knowledge_id": n.knowledge_id, + "knowledge_type": n.knowledge_type, + "priority": n.priority.value, + "requires_action": n.requires_action, + "status": n.status.value, + "suggested_actions": n.suggested_actions, + "metadata": n.metadata, + "created_at": n.created_at.isoformat(), + "expires_at": n.expires_at.isoformat() if n.expires_at else None, + } + for n in notifications + ], + "count": len(notifications), + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get knowledge notifications: {str(e)}" + ) + + +@app.put( + "/knowledge/notifications/{notification_id}/status", + tags=["knowledge-management"], + summary="Update Notification Status", + description="Mark notification as read, acknowledged, or acted upon", +) +async def update_notification_status( + notification_id: str = Path(..., description="Notification ID"), + status: str = Body(..., description="New notification status"), + action_taken: Optional[Dict[str, Any]] = Body( + None, description="Optional action taken metadata" + ), +): + """Update notification status and optional action taken""" + try: + from .knowledge_notification_service import NotificationStatus + + # Validate status + try: + notification_status = NotificationStatus(status) + except ValueError: + raise HTTPException(status_code=400, detail=f"Invalid status: {status}") + + success = await app.state.notification_service.mark_notification_status( + notification_id=notification_id, + status=notification_status, + action_taken=action_taken, + ) + + if success: + return { + "notification_id": notification_id, + "status": status, + "updated": True, + } + else: + raise HTTPException(status_code=404, detail="Notification not found") + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to update notification status: {str(e)}" + ) + + +@app.get( + "/knowledge/notifications/statistics", + tags=["knowledge-management"], + summary="Get Notification Statistics", + description="Get comprehensive notification statistics and analytics", +) +async def get_notification_statistics( + organization_id: Optional[str] = Query( + None, description="Filter by organization ID" + ), + days_back: int = Query(30, ge=1, le=365, description="Days of history to analyze"), +): + """Get notification statistics and analytics""" + try: + stats = await app.state.notification_service.get_notification_statistics( + organization_id=organization_id, days_back=days_back + ) + + return stats + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get notification statistics: {str(e)}" + ) + + +@app.post( + "/knowledge/organizations/{organization_id}/add", + tags=["knowledge-management"], + summary="Add Organizational Knowledge", + description="Add knowledge to organization-level knowledge base", +) +async def add_organizational_knowledge( + organization_id: str = Path(..., description="Organization ID"), + title: str = Body(..., description="Knowledge title"), + content: str = Body(..., description="Knowledge content"), + content_type: str = Body("documentation", description="Content type"), + knowledge_category: str = Body("development", description="Knowledge category"), + source_agent_id: Optional[str] = Body(None, description="Source agent ID"), + source_team_id: Optional[str] = Body(None, description="Source team ID"), + tags: List[str] = Body(default_factory=list, description="Knowledge tags"), + metadata: Dict[str, Any] = Body( + default_factory=dict, description="Additional metadata" + ), +): + """Add knowledge to organization-level knowledge base""" + try: + from .organization_rag_manager import ( + ContentType, + KnowledgeCategory, + OrganizationRAGManager, + SourceType, + ) + + # Initialize services if not already done + if not hasattr(app.state, "org_rag_manager"): + database_url = os.getenv( + "DATABASE_URL", + "postgresql://postgres:password@postgres:5432/ai_context", + ) + app.state.org_rag_manager = OrganizationRAGManager(database_url) + await app.state.org_rag_manager.initialize() + + # Validate enums + try: + content_type_enum = ContentType(content_type) + category_enum = KnowledgeCategory(knowledge_category) + except ValueError as e: + raise HTTPException(status_code=400, detail=f"Invalid enum value: {str(e)}") + + knowledge_id = await app.state.org_rag_manager.add_knowledge( + organization_id=organization_id, + title=title, + content=content, + content_type=content_type_enum, + knowledge_category=category_enum, + source_type=SourceType.MANUAL_INPUT, + source_agent_id=source_agent_id, + source_team_id=source_team_id, + tags=tags, + metadata=metadata, + ) + + return { + "knowledge_id": knowledge_id, + "organization_id": organization_id, + "title": title, + "status": "added", + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to add organizational knowledge: {str(e)}" + ) + + +@app.get( + "/knowledge/organizations/{organization_id}/search", + tags=["knowledge-management"], + summary="Search Organizational Knowledge", + description="Search organization-level knowledge base", +) +async def search_organizational_knowledge( + organization_id: str = Path(..., description="Organization ID"), + query: str = Query(..., description="Search query"), + limit: int = Query(10, ge=1, le=50, description="Maximum results to return"), + min_similarity: float = Query( + 0.3, ge=0.0, le=1.0, description="Minimum similarity threshold" + ), + categories: Optional[str] = Query( + None, description="Filter by categories (comma-separated)" + ), +): + """Search organization-level knowledge base""" + try: + from .organization_rag_manager import KnowledgeCategory + + # Parse categories + category_filters = None + if categories: + try: + category_filters = [ + KnowledgeCategory(cat.strip()) for cat in categories.split(",") + ] + except ValueError as e: + raise HTTPException( + status_code=400, detail=f"Invalid category: {str(e)}" + ) + + search_results = await app.state.org_rag_manager.search_knowledge( + organization_id=organization_id, + query=query, + limit=limit, + min_similarity=min_similarity, + categories=category_filters, + ) + + results = [] + for result in search_results: + results.append( + { + "knowledge_id": result.knowledge.id, + "title": result.knowledge.title, + "content_preview": ( + result.knowledge.content[:200] + "..." + if len(result.knowledge.content) > 200 + else result.knowledge.content + ), + "category": result.knowledge.knowledge_category.value, + "content_type": result.knowledge.content_type.value, + "similarity_score": result.similarity_score, + "combined_score": result.combined_score, + "quality_score": result.knowledge.quality_score, + "usage_count": result.knowledge.usage_count, + "created_at": result.knowledge.created_at.isoformat(), + "tags": result.knowledge.tags, + "metadata": result.knowledge.metadata, + } + ) + + return { + "organization_id": organization_id, + "query": query, + "results": results, + "count": len(results), + } + + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"Failed to search organizational knowledge: {str(e)}", + ) + + +@app.get( + "/knowledge/context-enhancement/{agent_id}", + tags=["knowledge-management"], + summary="Get Enhanced Context for Agent", + description="Get enhanced context with relevant organizational knowledge for task execution", +) +async def get_enhanced_context_for_agent( + agent_id: str = Path(..., description="Agent ID"), + task_description: str = Query( + ..., description="Task description for context enhancement" + ), + task_type: Optional[str] = Query(None, description="Task type"), + technologies: Optional[str] = Query( + None, description="Technologies involved (comma-separated)" + ), +): + """Get enhanced context with relevant knowledge for agent task execution""" + try: + from .context_enhancement_service import ContextEnhancementService + + # Initialize context enhancement service if needed + if not hasattr(app.state, "context_enhancement_service"): + database_url = os.getenv( + "DATABASE_URL", + "postgresql://postgres:password@postgres:5432/ai_context", + ) + # These would be initialized in the lifespan + if hasattr(app.state, "org_rag_manager") and hasattr( + app.state, "team_knowledge_manager" + ): + app.state.context_enhancement_service = ContextEnhancementService( + database_url=database_url, + org_rag_manager=app.state.org_rag_manager, + team_knowledge_manager=app.state.team_knowledge_manager, + ) + await app.state.context_enhancement_service.initialize() + else: + raise HTTPException( + status_code=503, + detail="Knowledge management services not initialized", + ) + + # Build task data + task_data = { + "description": task_description, + "task_type": task_type, + "technologies": technologies.split(",") if technologies else [], + } + + enhanced_context = ( + await app.state.context_enhancement_service.enhance_agent_context( + agent_id=agent_id, task_data=task_data + ) + ) + + return { + "agent_id": agent_id, + "task_description": task_description, + "enhanced_context": { + "organizational_knowledge_count": len( + enhanced_context.organizational_knowledge + ), + "team_knowledge_count": len(enhanced_context.team_knowledge), + "similar_task_insights_count": len( + enhanced_context.similar_task_insights + ), + "success_patterns": enhanced_context.success_patterns, + "common_pitfalls": enhanced_context.common_pitfalls, + "recommended_approaches": enhanced_context.recommended_approaches, + "context_summary": enhanced_context.context_summary, + "enhancement_metadata": enhanced_context.enhancement_metadata, + }, + "organizational_knowledge": [ + { + "knowledge_id": item.knowledge_id, + "title": item.title, + "category": item.category, + "relevance_score": item.relevance_score, + "confidence_score": item.confidence_score, + "content_preview": ( + item.content[:200] + "..." + if len(item.content) > 200 + else item.content + ), + } + for item in enhanced_context.organizational_knowledge + ], + "team_knowledge": [ + { + "knowledge_id": item.knowledge_id, + "title": item.title, + "category": item.category, + "relevance_score": item.relevance_score, + "confidence_score": item.confidence_score, + "content_preview": ( + item.content[:200] + "..." + if len(item.content) > 200 + else item.content + ), + } + for item in enhanced_context.team_knowledge + ], + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get enhanced context: {str(e)}" + ) + + +@app.get( + "/knowledge/analytics/organizations/{organization_id}/insights", + tags=["knowledge-management"], + summary="Get Organizational Knowledge Insights", + description="Get comprehensive analytics and insights about organizational knowledge", +) +async def get_organizational_knowledge_insights( + organization_id: str = Path(..., description="Organization ID"), + analysis_period_days: int = Query( + 30, ge=7, le=365, description="Analysis period in days" + ), +): + """Get comprehensive organizational knowledge insights and analytics""" + try: + insights = ( + await app.state.knowledge_analytics_service.get_organizational_insights( + organization_id=organization_id, + analysis_period_days=analysis_period_days, + ) + ) + + return { + "organization_id": organization_id, + "analysis_period_days": analysis_period_days, + "insights": { + "total_knowledge_items": insights.total_knowledge_items, + "knowledge_growth_rate": insights.knowledge_growth_rate, + "knowledge_utilization_rate": insights.knowledge_utilization_rate, + "knowledge_freshness_score": insights.knowledge_freshness_score, + "cross_team_sharing_rate": insights.cross_team_sharing_rate, + "propagation_efficiency": insights.propagation_efficiency, + "top_performing_categories": insights.top_performing_categories, + "knowledge_gaps": insights.knowledge_gaps, + "agent_knowledge_engagement": insights.agent_knowledge_engagement, + "team_knowledge_contribution": insights.team_knowledge_contribution, + "recommendations": insights.recommendations, + }, + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get organizational insights: {str(e)}" + ) + + +@app.get( + "/knowledge/analytics/organizations/{organization_id}/effectiveness", + tags=["knowledge-management"], + summary="Analyze Knowledge Effectiveness", + description="Analyze effectiveness and performance of knowledge items", +) +async def analyze_knowledge_effectiveness( + organization_id: str = Path(..., description="Organization ID"), + knowledge_category: Optional[str] = Query( + None, description="Filter by knowledge category" + ), + min_usage_count: int = Query( + 3, ge=1, description="Minimum usage count for analysis" + ), +): + """Analyze effectiveness of knowledge items in the organization""" + try: + effectiveness_metrics = ( + await app.state.knowledge_analytics_service.analyze_knowledge_effectiveness( + organization_id=organization_id, + knowledge_category=knowledge_category, + min_usage_count=min_usage_count, + ) + ) + + results = [] + for metric in effectiveness_metrics: + results.append( + { + "knowledge_id": metric.knowledge_id, + "title": metric.title, + "category": metric.category, + "usage_count": metric.usage_count, + "success_correlation": metric.success_correlation, + "average_relevance": metric.average_relevance, + "agent_adoption_rate": metric.agent_adoption_rate, + "team_adoption_rate": metric.team_adoption_rate, + "quality_score": metric.quality_score, + "recency_score": metric.recency_score, + "overall_effectiveness": metric.overall_effectiveness, + "trend_direction": metric.trend_direction, + "optimization_suggestions": metric.optimization_suggestions, + } + ) + + return { + "organization_id": organization_id, + "effectiveness_analysis": results, + "total_analyzed": len(results), + "summary": { + "avg_effectiveness": sum(r["overall_effectiveness"] for r in results) + / max(len(results), 1), + "top_performers": sorted( + results, key=lambda x: x["overall_effectiveness"], reverse=True + )[:5], + "needs_attention": [ + r for r in results if r["overall_effectiveness"] < 0.5 + ], + }, + } + + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"Failed to analyze knowledge effectiveness: {str(e)}", + ) + + +@app.get( + "/knowledge/analytics/agents/{agent_id}/profile", + tags=["knowledge-management"], + summary="Get Agent Knowledge Profile", + description="Get detailed knowledge profile and analytics for an agent", +) +async def get_agent_knowledge_profile( + agent_id: str = Path(..., description="Agent ID"), + analysis_period_days: int = Query( + 60, ge=7, le=365, description="Analysis period in days" + ), +): + """Get detailed knowledge profile for an agent""" + try: + profile = ( + await app.state.knowledge_analytics_service.get_agent_knowledge_profile( + agent_id=agent_id, analysis_period_days=analysis_period_days + ) + ) + + if not profile: + raise HTTPException( + status_code=404, detail="Agent not found or no knowledge data available" + ) + + return { + "agent_id": agent_id, + "analysis_period_days": analysis_period_days, + "profile": { + "agent_name": profile.agent_name, + "team_id": profile.team_id, + "knowledge_consumption_rate": profile.knowledge_consumption_rate, + "knowledge_creation_rate": profile.knowledge_creation_rate, + "expertise_areas": profile.expertise_areas, + "knowledge_application_success": profile.knowledge_application_success, + "learning_velocity": profile.learning_velocity, + "knowledge_sharing_activity": profile.knowledge_sharing_activity, + "preferred_knowledge_types": profile.preferred_knowledge_types, + "knowledge_gaps": profile.knowledge_gaps, + "optimization_recommendations": profile.optimization_recommendations, + }, + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get agent knowledge profile: {str(e)}" + ) + + +@app.get( + "/knowledge/analytics/organizations/{organization_id}/optimization", + tags=["knowledge-management"], + summary="Get Knowledge Optimization Recommendations", + description="Get comprehensive recommendations for knowledge system optimization", +) +async def get_knowledge_optimization_recommendations( + organization_id: str = Path(..., description="Organization ID"), + focus_area: Optional[str] = Query( + None, + description="Focus area (utilization, quality, gaps, propagation, collaboration)", + ), +): + """Generate comprehensive knowledge optimization recommendations""" + try: + recommendations = await app.state.knowledge_analytics_service.generate_knowledge_optimization_recommendations( + organization_id=organization_id, focus_area=focus_area + ) + + return { + "organization_id": organization_id, + "focus_area": focus_area, + "recommendations": recommendations, + "total_recommendations": len(recommendations), + } + + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"Failed to get optimization recommendations: {str(e)}", + ) + + +@app.get( + "/knowledge/analytics/organizations/{organization_id}/trends", + tags=["knowledge-management"], + summary="Get Knowledge Trends Analysis", + description="Analyze knowledge trends and patterns over time", +) +async def get_knowledge_trends_analysis( + organization_id: str = Path(..., description="Organization ID"), + trend_period_days: int = Query( + 90, ge=30, le=365, description="Trend analysis period in days" + ), +): + """Get comprehensive knowledge trends analysis""" + try: + trends = ( + await app.state.knowledge_analytics_service.get_knowledge_trends_analysis( + organization_id=organization_id, trend_period_days=trend_period_days + ) + ) + + return { + "organization_id": organization_id, + "trend_period_days": trend_period_days, + "trends": trends, + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get knowledge trends: {str(e)}" + ) + + +# Memory-Enhanced Agents Endpoints + + +@app.post( + "/agents/{agent_id}/deploy-memory", + tags=["memory-agents"], + summary="Deploy Memory-Enabled Agent", + description="Deploy an agent with persistent memory capabilities", +) +async def deploy_memory_enabled_agent( + agent_id: str = Path(..., description="Agent ID"), + template_id: str = Body(..., description="Agent template ID"), + task_id: Optional[str] = Body(None, description="Optional specific task ID"), + repository_settings: Optional[Dict[str, Any]] = Body( + None, description="Repository settings" + ), +): + """Deploy a memory-enabled autonomous agent container""" + try: + result = await app.state.agent_manager.deploy_memory_enabled_agent( + agent_id=agent_id, + template_id=template_id, + task_id=task_id, + repository_settings=repository_settings, + ) + + if result["success"]: + return result + else: + raise HTTPException(status_code=500, detail=result["error"]) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to deploy memory-enabled agent: {str(e)}" + ) + + +@app.get( + "/agents/{agent_id}/memory-status", + tags=["memory-agents"], + summary="Get Agent Memory Status", + description="Get agent memory status and expertise summary", +) +async def get_agent_memory_status(agent_id: str = Path(..., description="Agent ID")): + """Get agent memory status, expertise metrics, and insights""" + try: + status = await app.state.agent_manager.get_agent_memory_status(agent_id) + return status + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get agent memory status: {str(e)}" + ) + + +@app.post( + "/agents/{agent_id}/memory-tasks", + tags=["memory-agents"], + summary="Assign Task to Memory Agent", + description="Assign a task to a memory-enabled agent", +) +async def assign_task_to_memory_agent( + agent_id: str = Path(..., description="Agent ID"), + task_id: str = Body(..., description="Task ID"), + task_data: Dict[str, Any] = Body(..., description="Task data"), +): + """Assign a task to a memory-enabled agent for autonomous execution""" + try: + result = await app.state.agent_manager.assign_task_to_memory_agent( + agent_id=agent_id, task_id=task_id, task_data=task_data + ) + + if result["success"]: + return result + else: + raise HTTPException(status_code=400, detail=result["error"]) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to assign task to memory agent: {str(e)}" + ) + + +@app.delete( + "/agents/{agent_id}/memory", + tags=["memory-agents"], + summary="Stop Memory-Enabled Agent", + description="Stop a memory-enabled agent container", +) +async def stop_memory_enabled_agent(agent_id: str = Path(..., description="Agent ID")): + """Stop and clean up a memory-enabled agent container""" + try: + result = await app.state.agent_manager.stop_memory_enabled_agent(agent_id) + + if result["success"]: + return result + else: + raise HTTPException(status_code=400, detail=result["error"]) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to stop memory-enabled agent: {str(e)}" + ) + + +@app.get( + "/system/expertise-dashboard", + tags=["memory-agents"], + summary="Get System Expertise Dashboard", + description="Get system-wide expertise and memory analytics", +) +async def get_system_expertise_dashboard(): + """Get comprehensive dashboard of system expertise and memory analytics""" + try: + dashboard = await app.state.agent_manager.get_system_expertise_dashboard() + return dashboard + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get expertise dashboard: {str(e)}" + ) + + +@app.get( + "/agents/{agent_id}/tasks/pending", + tags=["memory-agents"], + summary="Get Pending Tasks for Agent", + description="Get pending tasks for a memory-enabled agent", +) +async def get_pending_tasks_for_agent( + agent_id: str = Path(..., description="Agent ID"), + limit: int = Query( + 10, ge=1, le=50, description="Maximum number of tasks to return" + ), +): + """Get pending tasks that a memory-enabled agent can pick up""" + try: + async with get_db_connection() as conn: + tasks = await conn.fetch( + """ + SELECT id, title, description, type, complexity, language, + framework, requirements, created_at + FROM tasks + WHERE agent_id = $1 + AND status = 'pending' + AND assigned_to_memory_agent = true + ORDER BY created_at ASC + LIMIT $2 + """, + agent_id, + limit, + ) + + return { + "agent_id": agent_id, + "tasks": [dict(task) for task in tasks], + "count": len(tasks), + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get pending tasks: {str(e)}" + ) + + +@app.put( + "/tasks/{task_id}/status", + tags=["memory-agents"], + summary="Update Task Status", + description="Update task status (used by memory-enabled agents)", +) +async def update_task_status( + task_id: str = Path(..., description="Task ID"), + status: str = Body(..., description="New task status"), + result: Optional[Dict[str, Any]] = Body(None, description="Task result data"), + updated_by: Optional[str] = Body(None, description="ID of agent updating the task"), + container_instance_id: Optional[str] = Body( + None, description="Container instance ID" + ), + updated_at: Optional[str] = Body(None, description="Update timestamp"), +): + """Update task status - used by memory-enabled agents to report progress""" + try: + async with get_db_connection() as conn: + await conn.execute( + """ + UPDATE tasks + SET status = $2, + result = COALESCE($3, result), + updated_by = COALESCE($4, updated_by), + updated_at = NOW() + WHERE id = $1 + """, + task_id, + status, + result, + updated_by, + ) + + # If task is completed, log it for expertise tracking + if status in ["completed", "failed"]: + # The agent's memory system will handle learning from the outcome + pass + + return {"task_id": task_id, "status": status, "updated": True} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to update task status: {str(e)}" + ) + + +@app.post( + "/agents/{agent_id}/register", + tags=["memory-agents"], + summary="Agent Registration", + description="Register agent capabilities and status with orchestrator", +) +async def register_agent_capabilities( + agent_id: str = Path(..., description="Agent ID"), + capabilities: Dict[str, Any] = Body( + ..., description="Agent capabilities and status" + ), +): + """Register or update agent capabilities - used by memory-enabled agents on startup""" + try: + # Update agent capabilities in database + async with get_db_connection() as conn: + await conn.execute( + """ + UPDATE agents + SET config = config || $2, + status = 'active', + updated_at = NOW() + WHERE id = $1 + """, + agent_id, + { + "capabilities": capabilities, + "last_registration": datetime.now().isoformat(), + }, + ) + + # Update in-memory tracking + if agent_id in app.state.agent_manager.memory_enabled_agents: + app.state.agent_manager.memory_enabled_agents[agent_id]["status"] = "active" + + return { + "agent_id": agent_id, + "agent_recognized": True, + "capabilities_accepted": True, + "status": "registered", + } + + except Exception as e: + return { + "agent_id": agent_id, + "agent_recognized": False, + "capabilities_accepted": False, + "error": str(e), + } + + +@app.post( + "/agents/{agent_id}/statistics", + tags=["memory-agents"], + summary="Agent Statistics Update", + description="Update agent performance and memory statistics", +) +async def update_agent_statistics( + agent_id: str = Path(..., description="Agent ID"), + stats: Dict[str, Any] = Body(..., description="Agent statistics"), +): + """Update agent statistics - used by memory-enabled agents for performance tracking""" + try: + # Store statistics for analytics + async with get_db_connection() as conn: + await conn.execute( + """ + UPDATE agents + SET config = config || $2, + updated_at = NOW() + WHERE id = $1 + """, + agent_id, + { + "latest_statistics": stats, + "statistics_updated_at": datetime.now().isoformat(), + }, + ) + + # Clear expertise cache to force refresh + await app.state.agent_manager.expertise_tracker.clear_cache(agent_id) + + return {"agent_id": agent_id, "statistics_updated": True} + + except Exception as e: + return {"agent_id": agent_id, "statistics_updated": False, "error": str(e)} + + +@app.post( + "/agents/{agent_id}/error", + tags=["memory-agents"], + summary="Agent Error Reporting", + description="Report agent errors for monitoring", +) +async def report_agent_error( + agent_id: str = Path(..., description="Agent ID"), + error_data: Dict[str, Any] = Body(..., description="Error information"), +): + """Report agent errors - used by memory-enabled agents for error tracking""" + try: + # Log error for monitoring + logger.error(f"Agent {agent_id} reported error: {error_data}") + + # Update agent status if it's a critical error + if error_data.get("critical", False): + async with get_db_connection() as conn: + await conn.execute( + """ + UPDATE agents + SET status = 'error', + config = config || $2, + updated_at = NOW() + WHERE id = $1 + """, + agent_id, + { + "last_error": error_data, + "error_reported_at": datetime.now().isoformat(), + }, + ) + + return {"agent_id": agent_id, "error_logged": True} + + except Exception as e: + logger.error(f"Failed to log agent error: {e}") + return {"agent_id": agent_id, "error_logged": False} + + +# ============================================================================ +# Goals Management API Endpoints +# ============================================================================ + + +@app.post( + "/organizations/{organization_id}/goals", + tags=["goals-management"], + summary="Create organizational goal", + description="Create a new goal for an organization with specified targets and deadlines", +) +async def create_goal( + organization_id: str = Path(..., description="Organization ID"), + goal_data: GoalCreateRequest = Body(..., description="Goal creation data"), + created_by: Optional[str] = Query( + None, description="ID of user/agent creating the goal" + ), +): + """Create a new organizational goal""" + try: + from .goals_management_service import GoalType + + goal_id = await app.state.goals_service.create_goal( + organization_id=organization_id, + title=goal_data.title, + description=goal_data.description, + goal_type=GoalType(goal_data.goal_type), + target_value=goal_data.target_value, + target_unit=goal_data.target_unit, + target_deadline=goal_data.target_deadline, + priority_level=goal_data.priority_level, + success_criteria=goal_data.success_criteria, + assigned_teams=goal_data.assigned_teams, + goal_owner_agent_id=goal_data.goal_owner_agent_id, + stakeholder_agents=goal_data.stakeholder_agents, + tags=goal_data.tags, + metadata=goal_data.metadata, + created_by=created_by, + ) + + return {"goal_id": goal_id, "status": "created"} + + except Exception as e: + logger.error(f"Error creating goal: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/organizations/{organization_id}/goals", + tags=["goals-management"], + summary="List organization goals", + description="Get all goals for an organization with optional filtering", +) +async def list_organization_goals( + organization_id: str = Path(..., description="Organization ID"), + status: Optional[List[str]] = Query(None, description="Filter by goal status"), + goal_type: Optional[List[str]] = Query(None, description="Filter by goal type"), + limit: int = Query( + 50, ge=1, le=100, description="Maximum number of goals to return" + ), +): + """List goals for an organization""" + try: + from .goals_management_service import GoalStatus, GoalType + + status_filter = [GoalStatus(s) for s in status] if status else None + type_filter = [GoalType(gt) for gt in goal_type] if goal_type else None + + goals = await app.state.goals_service.list_organization_goals( + organization_id=organization_id, + status_filter=status_filter, + goal_type_filter=type_filter, + limit=limit, + ) + + return { + "organization_id": organization_id, + "goals": [ + { + "id": goal.id, + "title": goal.title, + "description": goal.description, + "goal_type": goal.goal_type.value, + "status": goal.status.value, + "progress_percentage": float(goal.progress_percentage), + "target_value": ( + float(goal.target_value) if goal.target_value else None + ), + "target_unit": goal.target_unit, + "current_value": ( + float(goal.current_value) if goal.current_value else None + ), + "target_deadline": goal.target_deadline.isoformat(), + "priority_level": goal.priority_level, + "completion_confidence": float(goal.completion_confidence), + "created_at": goal.created_at.isoformat(), + "updated_at": goal.updated_at.isoformat(), + } + for goal in goals + ], + } + + except Exception as e: + logger.error(f"Error listing organization goals: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/goals/{goal_id}", + tags=["goals-management"], + summary="Get goal details", + description="Get detailed information about a specific goal", +) +async def get_goal(goal_id: str = Path(..., description="Goal ID")): + """Get goal details""" + try: + goal = await app.state.goals_service.get_goal(goal_id) + + if not goal: + raise HTTPException(status_code=404, detail="Goal not found") + + return { + "id": goal.id, + "organization_id": goal.organization_id, + "title": goal.title, + "description": goal.description, + "goal_type": goal.goal_type.value, + "status": goal.status.value, + "progress_percentage": float(goal.progress_percentage), + "target_value": float(goal.target_value) if goal.target_value else None, + "target_unit": goal.target_unit, + "current_value": float(goal.current_value) if goal.current_value else None, + "success_criteria": goal.success_criteria, + "start_date": goal.start_date.isoformat(), + "target_deadline": goal.target_deadline.isoformat(), + "actual_completion_date": ( + goal.actual_completion_date.isoformat() + if goal.actual_completion_date + else None + ), + "priority_level": goal.priority_level, + "completion_confidence": float(goal.completion_confidence), + "assigned_teams": goal.assigned_teams, + "goal_owner_agent_id": goal.goal_owner_agent_id, + "stakeholder_agents": goal.stakeholder_agents, + "tags": goal.tags, + "metadata": goal.metadata, + "created_by": goal.created_by, + "created_at": goal.created_at.isoformat(), + "updated_at": goal.updated_at.isoformat(), + } + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting goal {goal_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/goals/{goal_id}/overview", + tags=["goals-management"], + summary="Get goal overview", + description="Get comprehensive overview of goal with milestones, tasks, and progress", +) +async def get_goal_overview(goal_id: str = Path(..., description="Goal ID")): + """Get comprehensive goal overview""" + try: + overview = await app.state.goals_service.get_goal_overview(goal_id) + + if not overview: + raise HTTPException(status_code=404, detail="Goal not found") + + return overview + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting goal overview {goal_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.put( + "/goals/{goal_id}/progress", + tags=["goals-management"], + summary="Update goal progress", + description="Update progress for a specific goal", +) +async def update_goal_progress( + goal_id: str = Path(..., description="Goal ID"), + progress_data: GoalUpdateRequest = Body(..., description="Progress update data"), + recorded_by: Optional[str] = Query( + None, description="ID of user/agent recording progress" + ), +): + """Update goal progress""" + try: + success = await app.state.goals_service.update_goal_progress( + goal_id=goal_id, + progress_percentage=progress_data.progress_percentage, + current_value=progress_data.current_value, + completion_confidence=progress_data.completion_confidence, + progress_notes=progress_data.notes, + recorded_by=recorded_by, + ) + + if not success: + raise HTTPException( + status_code=404, detail="Goal not found or no changes made" + ) + + return {"goal_id": goal_id, "status": "updated"} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error updating goal progress {goal_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/goals/{goal_id}/milestones", + tags=["goals-management"], + summary="Create milestone", + description="Create a new milestone for a goal", +) +async def create_milestone( + goal_id: str = Path(..., description="Goal ID"), + milestone_data: MilestoneCreateRequest = Body( + ..., description="Milestone creation data" + ), + created_by: Optional[str] = Query( + None, description="ID of user/agent creating milestone" + ), +): + """Create milestone for goal""" + try: + milestone_id = await app.state.goals_service.create_milestone( + goal_id=goal_id, + title=milestone_data.title, + description=milestone_data.description, + target_date=milestone_data.target_date, + milestone_type=milestone_data.milestone_type, + success_criteria=milestone_data.success_criteria, + deliverables=milestone_data.deliverables, + dependencies=milestone_data.dependencies, + assigned_teams=milestone_data.assigned_teams, + responsible_agent_id=milestone_data.responsible_agent_id, + priority_level=milestone_data.priority_level, + weight_in_goal=milestone_data.weight_in_goal, + created_by=created_by, + ) + + return {"milestone_id": milestone_id, "status": "created"} + + except Exception as e: + logger.error(f"Error creating milestone: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/milestones/{milestone_id}/tasks", + tags=["goals-management"], + summary="Create task from milestone", + description="Create a new task derived from a milestone", +) +async def create_task_from_milestone( + milestone_id: str = Path(..., description="Milestone ID"), + task_data: TaskFromMilestoneRequest = Body(..., description="Task creation data"), + created_by: Optional[str] = Query( + None, description="ID of user/agent creating task" + ), +): + """Create task from milestone""" + try: + task_id = await app.state.goals_service.create_task_from_milestone( + milestone_id=milestone_id, + title=task_data.title, + description=task_data.description, + task_type=task_data.task_type, + complexity_level=task_data.complexity_level, + estimated_hours=task_data.estimated_hours, + due_date=task_data.due_date, + assigned_team_id=task_data.assigned_team_id, + assigned_agent_id=task_data.assigned_agent_id, + priority=task_data.priority, + requirements=task_data.requirements, + acceptance_criteria=task_data.acceptance_criteria, + dependencies=task_data.dependencies, + created_by_agent_id=created_by, + ) + + return {"task_id": task_id, "status": "created"} + + except Exception as e: + logger.error(f"Error creating task from milestone: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/goals/{goal_id}/generate-execution-plan", + tags=["goals-management"], + summary="Generate execution plan", + description="Generate comprehensive milestone and task execution plan for a goal", +) +async def generate_execution_plan( + goal_id: str = Path(..., description="Goal ID"), + planning_context: Optional[Dict[str, Any]] = Body( + None, description="Additional planning context" + ), +): + """Generate execution plan with milestones and tasks""" + try: + execution_plan = ( + await app.state.milestone_task_engine.generate_goal_execution_plan( + goal_id=goal_id, planning_context=planning_context + ) + ) + + return execution_plan + + except Exception as e: + logger.error(f"Error generating execution plan for goal {goal_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/goals/{goal_id}/generate-monthly-milestones", + tags=["goals-management"], + summary="Generate monthly milestones", + description="Generate monthly milestone breakdown for a goal", +) +async def generate_monthly_milestones( + goal_id: str = Path(..., description="Goal ID"), + start_date: Optional[date] = Query(None, description="Start date for milestones"), + end_date: Optional[date] = Query(None, description="End date for milestones"), +): + """Generate monthly milestones for goal""" + try: + milestone_ids = ( + await app.state.milestone_task_engine.generate_monthly_milestones( + goal_id=goal_id, start_date=start_date, end_date=end_date + ) + ) + + return { + "goal_id": goal_id, + "milestone_ids": milestone_ids, + "count": len(milestone_ids), + "status": "generated", + } + + except Exception as e: + logger.error(f"Error generating monthly milestones for goal {goal_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/milestones/{milestone_id}/generate-weekly-tasks", + tags=["goals-management"], + summary="Generate weekly tasks", + description="Generate weekly task breakdown for a milestone", +) +async def generate_weekly_tasks( + milestone_id: str = Path(..., description="Milestone ID"), + focus_areas: Optional[List[str]] = Body( + None, description="Focus areas for task generation" + ), +): + """Generate weekly tasks for milestone""" + try: + task_ids = ( + await app.state.milestone_task_engine.generate_weekly_tasks_for_milestone( + milestone_id=milestone_id, focus_areas=focus_areas + ) + ) + + return { + "milestone_id": milestone_id, + "task_ids": task_ids, + "count": len(task_ids), + "status": "generated", + } + + except Exception as e: + logger.error(f"Error generating weekly tasks for milestone {milestone_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/goals/{goal_id}/generate-cross-functional-tasks", + tags=["goals-management"], + summary="Generate cross-functional tasks", + description="Generate tasks across different business functions for a goal", +) +async def generate_cross_functional_tasks( + goal_id: str = Path(..., description="Goal ID"), + target_functions: Optional[List[str]] = Body( + None, description="Target business functions" + ), +): + """Generate cross-functional tasks for goal""" + try: + functional_tasks = ( + await app.state.milestone_task_engine.generate_cross_functional_tasks( + goal_id=goal_id, target_functions=target_functions + ) + ) + + return { + "goal_id": goal_id, + "functional_tasks": functional_tasks, + "total_tasks": sum(len(tasks) for tasks in functional_tasks.values()), + "status": "generated", + } + + except Exception as e: + logger.error(f"Error generating cross-functional tasks for goal {goal_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/goals/{goal_id}/conversations", + tags=["goals-management"], + summary="Create goal conversation", + description="Create AI-powered conversation for goal planning and discussion", +) +async def create_goal_conversation( + goal_id: str = Path(..., description="Goal ID"), + conversation_data: GoalConversationCreateRequest = Body( + ..., description="Conversation creation data" + ), + created_by: Optional[str] = Query( + None, description="ID of user/agent creating conversation" + ), +): + """Create goal conversation""" + try: + from .goal_conversation_service import ConversationType + + conversation_id = ( + await app.state.goal_conversation_service.create_goal_conversation( + goal_id=goal_id, + conversation_type=ConversationType(conversation_data.conversation_type), + conversation_title=conversation_data.conversation_title, + initial_context=conversation_data.initial_context, + participants=conversation_data.participants, + created_by=created_by, + ) + ) + + return {"conversation_id": conversation_id, "status": "created"} + + except Exception as e: + logger.error(f"Error creating goal conversation: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/conversations/{conversation_id}", + tags=["goals-management"], + summary="Get goal conversation", + description="Get full conversation with messages, insights, and action items", +) +async def get_goal_conversation( + conversation_id: str = Path(..., description="Conversation ID") +): + """Get goal conversation""" + try: + conversation = await app.state.goal_conversation_service.get_conversation( + conversation_id + ) + + if not conversation: + raise HTTPException(status_code=404, detail="Conversation not found") + + return conversation + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting conversation {conversation_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/conversations/{conversation_id}/messages", + tags=["goals-management"], + summary="Add message to conversation", + description="Add a new message to a goal conversation", +) +async def add_message_to_conversation( + conversation_id: str = Path(..., description="Conversation ID"), + message_data: ConversationMessageRequest = Body(..., description="Message data"), + sender_id: Optional[str] = Query(None, description="ID of message sender"), +): + """Add message to conversation""" + try: + from .goal_conversation_service import MessageType + + message_id = ( + await app.state.goal_conversation_service.add_message_to_conversation( + conversation_id=conversation_id, + message_type=MessageType(message_data.message_type), + sender_id=sender_id, + sender_name=message_data.sender_name, + content=message_data.content, + metadata=message_data.metadata, + references=message_data.references, + ) + ) + + return {"message_id": message_id, "status": "added"} + + except Exception as e: + logger.error(f"Error adding message to conversation: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/conversations/{conversation_id}/generate-milestones", + tags=["goals-management"], + summary="Generate milestones from conversation", + description="Generate milestone recommendations based on conversation analysis", +) +async def generate_planning_milestones( + conversation_id: str = Path(..., description="Conversation ID"), + planning_context: Optional[Dict[str, Any]] = Body( + None, description="Additional planning context" + ), +): + """Generate planning milestones from conversation""" + try: + milestones = ( + await app.state.goal_conversation_service.generate_planning_milestones( + conversation_id=conversation_id, planning_context=planning_context + ) + ) + + return { + "conversation_id": conversation_id, + "milestones": milestones, + "count": len(milestones), + "status": "generated", + } + + except Exception as e: + logger.error(f"Error generating planning milestones: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/conversations/{conversation_id}/conduct-progress-review", + tags=["goals-management"], + summary="Conduct progress review", + description="Conduct AI-powered progress review for a goal conversation", +) +async def conduct_progress_review( + conversation_id: str = Path(..., description="Conversation ID"), + review_period_days: int = Query( + 30, ge=1, le=365, description="Review period in days" + ), +): + """Conduct progress review""" + try: + review_analysis = ( + await app.state.goal_conversation_service.conduct_progress_review( + conversation_id=conversation_id, review_period_days=review_period_days + ) + ) + + return review_analysis + + except Exception as e: + logger.error(f"Error conducting progress review: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/conversations/{conversation_id}/extract-action-items", + tags=["goals-management"], + summary="Extract action items", + description="Extract and create action items from conversation analysis", +) +async def extract_action_items( + conversation_id: str = Path(..., description="Conversation ID"), + auto_assign: bool = Query( + True, description="Whether to automatically assign action items" + ), +): + """Extract action items from conversation""" + try: + action_items = await app.state.goal_conversation_service.extract_action_items_from_conversation( + conversation_id=conversation_id, auto_assign=auto_assign + ) + + return { + "conversation_id": conversation_id, + "action_items": action_items, + "count": len(action_items), + "status": "extracted", + } + + except Exception as e: + logger.error(f"Error extracting action items: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/goals/{goal_id}/conversations", + tags=["goals-management"], + summary="Get goal conversations", + description="Get all conversations for a goal with optional filtering", +) +async def get_goal_conversations( + goal_id: str = Path(..., description="Goal ID"), + conversation_type: Optional[str] = Query( + None, description="Filter by conversation type" + ), + status: Optional[str] = Query(None, description="Filter by conversation status"), + limit: int = Query(10, ge=1, le=50, description="Maximum number of conversations"), +): + """Get conversations for a goal""" + try: + from .goal_conversation_service import ConversationStatus, ConversationType + + conv_type = ConversationType(conversation_type) if conversation_type else None + conv_status = ConversationStatus(status) if status else None + + conversations = ( + await app.state.goal_conversation_service.get_goal_conversations( + goal_id=goal_id, + conversation_type=conv_type, + status=conv_status, + limit=limit, + ) + ) + + return { + "goal_id": goal_id, + "conversations": conversations, + "count": len(conversations), + } + + except Exception as e: + logger.error(f"Error getting goal conversations: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/goals/{goal_id}/track-progress", + tags=["goals-management"], + summary="Record progress tracking update", + description="Record detailed progress update with tracking and risk assessment", +) +async def record_progress_tracking( + goal_id: str = Path(..., description="Goal ID"), + progress_data: ProgressUpdateRequest = Body( + ..., description="Progress tracking data" + ), + recorded_by: Optional[str] = Query( + None, description="ID of user/agent recording progress" + ), +): + """Record progress tracking update""" + try: + snapshot_id = await app.state.goal_tracking_service.record_progress_update( + goal_id=goal_id, + progress_percentage=progress_data.progress_percentage, + current_value=progress_data.current_value, + milestone_id=progress_data.milestone_id, + notes=progress_data.notes, + recorded_by=recorded_by, + confidence_score=progress_data.confidence_score, + trigger_alerts=progress_data.trigger_alerts, + ) + + return {"goal_id": goal_id, "snapshot_id": snapshot_id, "status": "recorded"} + + except Exception as e: + logger.error(f"Error recording progress tracking: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/goals/{goal_id}/deadline-risk", + tags=["goals-management"], + summary="Assess deadline risk", + description="Get comprehensive deadline risk assessment for a goal", +) +async def assess_deadline_risk(goal_id: str = Path(..., description="Goal ID")): + """Assess deadline risk for goal""" + try: + deadline_risk = await app.state.goal_tracking_service.assess_goal_deadline_risk( + goal_id + ) + + return { + "goal_id": deadline_risk.goal_id, + "risk_level": deadline_risk.risk_level.value, + "probability_of_delay": float(deadline_risk.probability_of_delay), + "estimated_completion_date": deadline_risk.estimated_completion_date.isoformat(), + "days_at_risk": deadline_risk.days_at_risk, + "critical_path_items": deadline_risk.critical_path_items, + "mitigation_strategies": deadline_risk.mitigation_strategies, + "updated_at": deadline_risk.updated_at.isoformat(), + } + + except Exception as e: + logger.error(f"Error assessing deadline risk: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/goals/{goal_id}/progress-report", + tags=["goals-management"], + summary="Generate progress report", + description="Generate comprehensive progress report for a goal", +) +async def generate_progress_report( + goal_id: str = Path(..., description="Goal ID"), + report_period_days: int = Query( + 30, ge=1, le=365, description="Report period in days" + ), +): + """Generate progress report for goal""" + try: + report = await app.state.goal_tracking_service.generate_progress_report( + goal_id=goal_id, report_period_days=report_period_days + ) + + return report + + except Exception as e: + logger.error(f"Error generating progress report: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/organizations/{organization_id}/goals-dashboard", + tags=["goals-management"], + summary="Get organization goals dashboard", + description="Get comprehensive dashboard for all organization goals", +) +async def get_organization_goals_dashboard( + organization_id: str = Path(..., description="Organization ID") +): + """Get organization goals dashboard""" + try: + dashboard = await app.state.goals_service.get_organization_goals_dashboard( + organization_id + ) + return dashboard + + except Exception as e: + logger.error(f"Error getting organization dashboard: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/organizations/{organization_id}/tracking-dashboard", + tags=["goals-management"], + summary="Get tracking dashboard", + description="Get comprehensive tracking dashboard with risk assessments", +) +async def get_tracking_dashboard( + organization_id: str = Path(..., description="Organization ID") +): + """Get organization tracking dashboard""" + try: + dashboard = ( + await app.state.goal_tracking_service.get_organization_tracking_dashboard( + organization_id + ) + ) + return dashboard + + except Exception as e: + logger.error(f"Error getting tracking dashboard: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# ================================ +# Knowledge Management API Endpoints +# ================================ + + +@app.post( + "/knowledge/organizations/{organization_id}/documents", + tags=["knowledge-management"], + summary="Upload Organizational Document", + response_model=DocumentMetadata, +) +async def upload_organization_document( + organization_id: str = Path(..., description="Organization ID"), + file: UploadFile = File(..., description="Document file to upload"), + title: Optional[str] = Form(None, description="Document title"), + tags: Optional[str] = Form(None, description="Comma-separated tags"), +): + """Upload a document to organizational knowledge base""" + try: + tags_list = [] + if tags: + tags_list = [tag.strip() for tag in tags.split(",")] + + document = await knowledge_manager.upload_document( + file_content=file.file, + filename=file.filename, + title=title, + organization_id=organization_id, + tags=tags_list, + ) + + return document + + except Exception as e: + logger.error(f"Error uploading organizational document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/knowledge/organizations/{organization_id}/url", + tags=["knowledge-management"], + summary="Add URL to Organizational Knowledge", + response_model=DocumentMetadata, +) +async def add_organization_url( + organization_id: str = Path(..., description="Organization ID"), + url: str = Body(..., embed=True), + title: Optional[str] = Body(None, embed=True), + tags: Optional[List[str]] = Body(None, embed=True), +): + """Add URL content to organizational knowledge base""" + try: + document = await knowledge_manager.upload_url( + url=url, title=title, organization_id=organization_id, tags=tags or [] + ) + + return document + + except Exception as e: + logger.error(f"Error adding organizational URL: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/organizations/{organization_id}/documents", + tags=["knowledge-management"], + summary="List Organizational Documents", + response_model=List[DocumentMetadata], +) +async def list_organization_documents( + organization_id: str = Path(..., description="Organization ID") +): + """Get list of organizational documents""" + try: + documents = await knowledge_manager.get_documents( + organization_id=organization_id + ) + return documents + + except Exception as e: + logger.error(f"Error listing organizational documents: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/organizations/{organization_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Get Organizational Document", + response_model=DocumentMetadata, +) +async def get_organization_document( + organization_id: str = Path(..., description="Organization ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Get organizational document metadata""" + try: + document = await knowledge_manager.get_document_metadata( + doc_id=doc_id, organization_id=organization_id + ) + + if not document: + raise HTTPException(status_code=404, detail="Document not found") + + return document + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting organizational document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/organizations/{organization_id}/documents/{doc_id}/content", + tags=["knowledge-management"], + summary="Get Organizational Document Content", +) +async def get_organization_document_content( + organization_id: str = Path(..., description="Organization ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Get full content of organizational document""" + try: + content = await knowledge_manager.get_document_content( + doc_id=doc_id, organization_id=organization_id + ) + + if content is None: + raise HTTPException(status_code=404, detail="Document not found") + + return {"content": content} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting organizational document content: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.put( + "/knowledge/organizations/{organization_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Update Organizational Document", + response_model=DocumentMetadata, +) +async def update_organization_document( + organization_id: str = Path(..., description="Organization ID"), + doc_id: str = Path(..., description="Document ID"), + title: Optional[str] = Body(None, embed=True), + tags: Optional[List[str]] = Body(None, embed=True), +): + """Update organizational document metadata""" + try: + document = await knowledge_manager.update_document( + doc_id=doc_id, title=title, tags=tags, organization_id=organization_id + ) + + if not document: + raise HTTPException(status_code=404, detail="Document not found") + + return document + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error updating organizational document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.delete( + "/knowledge/organizations/{organization_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Delete Organizational Document", +) +async def delete_organization_document( + organization_id: str = Path(..., description="Organization ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Delete organizational document""" + try: + success = await knowledge_manager.delete_document( + doc_id=doc_id, organization_id=organization_id + ) + + if not success: + raise HTTPException(status_code=404, detail="Document not found") + + return {"message": "Document deleted successfully"} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error deleting organizational document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# Team Knowledge Management Endpoints + + +@app.post( + "/knowledge/teams/{team_id}/documents", + tags=["knowledge-management"], + summary="Upload Team Document", + response_model=DocumentMetadata, +) +async def upload_team_document( + team_id: str = Path(..., description="Team ID"), + file: UploadFile = File(..., description="Document file to upload"), + title: Optional[str] = Form(None, description="Document title"), + tags: Optional[str] = Form(None, description="Comma-separated tags"), +): + """Upload a document to team knowledge base""" + try: + tags_list = [] + if tags: + tags_list = [tag.strip() for tag in tags.split(",")] + + document = await knowledge_manager.upload_document( + file_content=file.file, + filename=file.filename, + title=title, + team_id=team_id, + tags=tags_list, + ) + + return document + + except Exception as e: + logger.error(f"Error uploading team document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/knowledge/teams/{team_id}/url", + tags=["knowledge-management"], + summary="Add URL to Team Knowledge", + response_model=DocumentMetadata, +) +async def add_team_url( + team_id: str = Path(..., description="Team ID"), + url: str = Body(..., embed=True), + title: Optional[str] = Body(None, embed=True), + tags: Optional[List[str]] = Body(None, embed=True), +): + """Add URL content to team knowledge base""" + try: + document = await knowledge_manager.upload_url( + url=url, title=title, team_id=team_id, tags=tags or [] + ) + + return document + + except Exception as e: + logger.error(f"Error adding team URL: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/teams/{team_id}/documents", + tags=["knowledge-management"], + summary="List Team Documents", + response_model=List[DocumentMetadata], +) +async def list_team_documents(team_id: str = Path(..., description="Team ID")): + """Get list of team documents""" + try: + documents = await knowledge_manager.get_documents(team_id=team_id) + return documents + + except Exception as e: + logger.error(f"Error listing team documents: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/teams/{team_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Get Team Document", + response_model=DocumentMetadata, +) +async def get_team_document( + team_id: str = Path(..., description="Team ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Get team document metadata""" + try: + document = await knowledge_manager.get_document_metadata( + doc_id=doc_id, team_id=team_id + ) + + if not document: + raise HTTPException(status_code=404, detail="Document not found") + + return document + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting team document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/teams/{team_id}/documents/{doc_id}/content", + tags=["knowledge-management"], + summary="Get Team Document Content", +) +async def get_team_document_content( + team_id: str = Path(..., description="Team ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Get full content of team document""" + try: + content = await knowledge_manager.get_document_content( + doc_id=doc_id, team_id=team_id + ) + + if content is None: + raise HTTPException(status_code=404, detail="Document not found") + + return {"content": content} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting team document content: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.put( + "/knowledge/teams/{team_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Update Team Document", + response_model=DocumentMetadata, +) +async def update_team_document( + team_id: str = Path(..., description="Team ID"), + doc_id: str = Path(..., description="Document ID"), + title: Optional[str] = Body(None, embed=True), + tags: Optional[List[str]] = Body(None, embed=True), +): + """Update team document metadata""" + try: + document = await knowledge_manager.update_document( + doc_id=doc_id, title=title, tags=tags, team_id=team_id + ) + + if not document: + raise HTTPException(status_code=404, detail="Document not found") + + return document + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error updating team document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.delete( + "/knowledge/teams/{team_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Delete Team Document", +) +async def delete_team_document( + team_id: str = Path(..., description="Team ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Delete team document""" + try: + success = await knowledge_manager.delete_document( + doc_id=doc_id, team_id=team_id + ) + + if not success: + raise HTTPException(status_code=404, detail="Document not found") + + return {"message": "Document deleted successfully"} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error deleting team document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# Agent Knowledge Management Endpoints + + +@app.post( + "/knowledge/agents/{agent_id}/documents", + tags=["knowledge-management"], + summary="Upload Agent Document", + response_model=DocumentMetadata, +) +async def upload_agent_document( + agent_id: str = Path(..., description="Agent ID"), + file: UploadFile = File(..., description="Document file to upload"), + title: Optional[str] = Form(None, description="Document title"), + tags: Optional[str] = Form(None, description="Comma-separated tags"), +): + """Upload a document to agent knowledge base""" + try: + tags_list = [] + if tags: + tags_list = [tag.strip() for tag in tags.split(",")] + + document = await knowledge_manager.upload_document( + file_content=file.file, + filename=file.filename, + title=title, + agent_id=agent_id, + tags=tags_list, + ) + + return document + + except Exception as e: + logger.error(f"Error uploading agent document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/knowledge/agents/{agent_id}/url", + tags=["knowledge-management"], + summary="Add URL to Agent Knowledge", + response_model=DocumentMetadata, +) +async def add_agent_url( + agent_id: str = Path(..., description="Agent ID"), + url: str = Body(..., embed=True), + title: Optional[str] = Body(None, embed=True), + tags: Optional[List[str]] = Body(None, embed=True), +): + """Add URL content to agent knowledge base""" + try: + document = await knowledge_manager.upload_url( + url=url, title=title, agent_id=agent_id, tags=tags or [] + ) + + return document + + except Exception as e: + logger.error(f"Error adding agent URL: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/agents/{agent_id}/documents", + tags=["knowledge-management"], + summary="List Agent Documents", + response_model=List[DocumentMetadata], +) +async def list_agent_documents(agent_id: str = Path(..., description="Agent ID")): + """Get list of agent documents""" + try: + documents = await knowledge_manager.get_documents(agent_id=agent_id) + return documents + + except Exception as e: + logger.error(f"Error listing agent documents: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/agents/{agent_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Get Agent Document", + response_model=DocumentMetadata, +) +async def get_agent_document( + agent_id: str = Path(..., description="Agent ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Get agent document metadata""" + try: + document = await knowledge_manager.get_document_metadata( + doc_id=doc_id, agent_id=agent_id + ) + + if not document: + raise HTTPException(status_code=404, detail="Document not found") + + return document + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting agent document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/agents/{agent_id}/documents/{doc_id}/content", + tags=["knowledge-management"], + summary="Get Agent Document Content", +) +async def get_agent_document_content( + agent_id: str = Path(..., description="Agent ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Get full content of agent document""" + try: + content = await knowledge_manager.get_document_content( + doc_id=doc_id, agent_id=agent_id + ) + + if content is None: + raise HTTPException(status_code=404, detail="Document not found") + + return {"content": content} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting agent document content: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.put( + "/knowledge/agents/{agent_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Update Agent Document", + response_model=DocumentMetadata, +) +async def update_agent_document( + agent_id: str = Path(..., description="Agent ID"), + doc_id: str = Path(..., description="Document ID"), + title: Optional[str] = Body(None, embed=True), + tags: Optional[List[str]] = Body(None, embed=True), +): + """Update agent document metadata""" + try: + document = await knowledge_manager.update_document( + doc_id=doc_id, title=title, tags=tags, agent_id=agent_id + ) + + if not document: + raise HTTPException(status_code=404, detail="Document not found") + + return document + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error updating agent document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.delete( + "/knowledge/agents/{agent_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Delete Agent Document", +) +async def delete_agent_document( + agent_id: str = Path(..., description="Agent ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Delete agent document""" + try: + success = await knowledge_manager.delete_document( + doc_id=doc_id, agent_id=agent_id + ) + + if not success: + raise HTTPException(status_code=404, detail="Document not found") + + return {"message": "Document deleted successfully"} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error deleting agent document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# Knowledge Search Endpoints + + +@app.get( + "/knowledge/search", + tags=["knowledge-management"], + summary="Search Knowledge Base", + response_model=List[DocumentMetadata], +) +async def search_knowledge( + query: str = Query(..., description="Search query"), + organization_id: Optional[str] = Query(None, description="Filter by organization"), + team_id: Optional[str] = Query(None, description="Filter by team"), + agent_id: Optional[str] = Query(None, description="Filter by agent"), + limit: int = Query(10, ge=1, le=100, description="Maximum number of results"), +): + """Search across knowledge base""" + try: + documents = await knowledge_manager.search_documents( + query=query, + organization_id=organization_id, + team_id=team_id, + agent_id=agent_id, + limit=limit, + ) + + return documents + + except Exception as e: + logger.error(f"Error searching knowledge: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# ================================ +# Container Management API Endpoints +# ================================ + + +@app.post( + "/agents/{agent_id}/container/create", + tags=["container-management"], + summary="Create Agent Container", + response_model=ContainerStatus, +) +async def create_agent_container( + agent_id: str = Path(..., description="Agent ID"), + config: Optional[ContainerConfig] = Body( + None, description="Container configuration" + ), +): + """Create a new container for an AI agent""" + try: + status = await container_manager.create_agent_container(agent_id, config) + return status + + except Exception as e: + logger.error(f"Error creating container for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/agents/{agent_id}/container/start", + tags=["container-management"], + summary="Start Agent Container", + response_model=ContainerStatus, +) +async def start_agent_container(agent_id: str = Path(..., description="Agent ID")): + """Start an agent container""" + try: + status = await container_manager.start_container(agent_id) + return status + + except RuntimeError as e: + raise HTTPException(status_code=404, detail=str(e)) + except Exception as e: + logger.error(f"Error starting container for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/agents/{agent_id}/container/stop", + tags=["container-management"], + summary="Stop Agent Container", + response_model=ContainerStatus, +) +async def stop_agent_container( + agent_id: str = Path(..., description="Agent ID"), + timeout: int = Body(30, description="Stop timeout in seconds"), +): + """Stop an agent container""" + try: + status = await container_manager.stop_container(agent_id, timeout) + return status + + except RuntimeError as e: + raise HTTPException(status_code=404, detail=str(e)) + except Exception as e: + logger.error(f"Error stopping container for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/agents/{agent_id}/container/restart", + tags=["container-management"], + summary="Restart Agent Container", + response_model=ContainerStatus, +) +async def restart_agent_container( + agent_id: str = Path(..., description="Agent ID"), + timeout: int = Body(30, description="Restart timeout in seconds"), +): + """Restart an agent container""" + try: + status = await container_manager.restart_container(agent_id, timeout) + return status + + except RuntimeError as e: + raise HTTPException(status_code=404, detail=str(e)) + except Exception as e: + logger.error(f"Error restarting container for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.delete( + "/agents/{agent_id}/container", + tags=["container-management"], + summary="Remove Agent Container", +) +async def remove_agent_container( + agent_id: str = Path(..., description="Agent ID"), + force: bool = Query(False, description="Force removal of running container"), +): + """Remove an agent container""" + try: + success = await container_manager.remove_container(agent_id, force) + + if success: + return {"message": f"Container for agent {agent_id} removed successfully"} + else: + raise HTTPException(status_code=500, detail="Failed to remove container") + + except Exception as e: + logger.error(f"Error removing container for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/agents/{agent_id}/container/status", + tags=["container-management"], + summary="Get Agent Container Status", + response_model=Optional[ContainerStatus], +) +async def get_agent_container_status(agent_id: str = Path(..., description="Agent ID")): + """Get container status for an agent""" + try: + status = await container_manager.get_container_status(agent_id) + return status + + except Exception as e: + logger.error(f"Error getting container status for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/containers/agents", + tags=["container-management"], + summary="List Agent Containers", + response_model=List[ContainerStatus], +) +async def list_agent_containers(): + """List all agent containers""" + try: + containers = await container_manager.list_agent_containers() + return containers + + except Exception as e: + logger.error(f"Error listing agent containers: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/agents/{agent_id}/container/logs", + tags=["container-management"], + summary="Get Agent Container Logs", +) +async def get_agent_container_logs( + agent_id: str = Path(..., description="Agent ID"), + tail: int = Query(100, ge=1, le=10000, description="Number of log lines to return"), + since: Optional[str] = Query( + None, description="Show logs since timestamp (ISO format)" + ), +): + """Get container logs for an agent""" + try: + since_dt = None + if since: + try: + since_dt = datetime.fromisoformat(since.replace("Z", "+00:00")) + except ValueError: + raise HTTPException(status_code=400, detail="Invalid timestamp format") + + logs = await container_manager.get_container_logs( + agent_id=agent_id, tail=tail, since=since_dt + ) + + return {"logs": logs} + + except RuntimeError as e: + raise HTTPException(status_code=404, detail=str(e)) + except Exception as e: + logger.error(f"Error getting container logs for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/agents/{agent_id}/container/execute", + tags=["container-management"], + summary="Execute Command in Container", +) +async def execute_container_command( + agent_id: str = Path(..., description="Agent ID"), + command: str = Body(..., description="Command to execute"), + working_dir: Optional[str] = Body(None, description="Working directory"), +): + """Execute a command in the agent container""" + try: + result = await container_manager.execute_command( + agent_id=agent_id, command=command, working_dir=working_dir + ) + + return result + + except RuntimeError as e: + raise HTTPException(status_code=404, detail=str(e)) + except Exception as e: + logger.error(f"Error executing command in container for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# WebSocket endpoint for real-time log streaming +@app.websocket("/agents/{agent_id}/container/logs/stream") +async def stream_agent_container_logs(websocket: WebSocket, agent_id: str): + """Stream container logs in real-time via WebSocket""" + await websocket.accept() + + try: + # Check if container exists + status = await container_manager.get_container_status(agent_id) + if not status: + await websocket.send_json({"error": "Container not found"}) + await websocket.close() + return + + await websocket.send_json({"status": "connected", "agent_id": agent_id}) + + # Stream logs + async for log_entry in container_manager.stream_container_logs(agent_id): + await websocket.send_json( + { + "timestamp": log_entry.timestamp.isoformat(), + "stream": log_entry.stream, + "message": log_entry.message, + } + ) + + except Exception as e: + logger.error(f"Error in log stream for agent {agent_id}: {e}") + try: + await websocket.send_json({"error": str(e)}) + except Exception: + logger.debug("Failed to send error frame on closing websocket") + finally: + try: + await websocket.close() + except Exception: + logger.debug("Failed to close websocket cleanly") + + +# ============================================================================ +# RAG (Retrieval-Augmented Generation) Endpoints +# ============================================================================ + + +@app.post("/rag/search") +async def search_knowledge_context( + query: str = Body(..., embed=True), + organization_id: Optional[str] = Body(None, embed=True), + team_id: Optional[str] = Body(None, embed=True), + agent_id: Optional[str] = Body(None, embed=True), + max_results: int = Body(5, embed=True), + similarity_threshold: float = Body(0.7, embed=True), +): + """Search for relevant knowledge context using RAG""" + try: + context = await rag_system.search_relevant_context( + query=query, + organization_id=organization_id, + team_id=team_id, + agent_id=agent_id, + max_results=max_results, + similarity_threshold=similarity_threshold, + ) + + return { + "query": context.query, + "relevant_chunks": [ + { + "document_id": chunk.document_id, + "document_title": chunk.metadata.get("document_title", "Unknown"), + "content": chunk.content, + "chunk_index": chunk.chunk_index, + "metadata": chunk.metadata, + } + for chunk in context.relevant_chunks + ], + "similarity_scores": context.similarity_scores, + "total_documents": context.total_documents, + "context_length": context.context_length, + } + + except Exception as e: + logger.error(f"Error searching knowledge context: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post("/rag/enhance-prompt") +async def enhance_prompt_with_context( + message: str = Body(..., embed=True), + organization_id: Optional[str] = Body(None, embed=True), + team_id: Optional[str] = Body(None, embed=True), + agent_id: Optional[str] = Body(None, embed=True), + max_context_length: int = Body(4000, embed=True), +): + """Enhance a prompt with relevant context using RAG""" + try: + enhanced_prompt = await rag_system.get_contextual_prompt( + user_message=message, + organization_id=organization_id, + team_id=team_id, + agent_id=agent_id, + max_context_length=max_context_length, + ) + + return { + "original_message": message, + "enhanced_prompt": enhanced_prompt, + "context_added": len(enhanced_prompt) > len(message), + } + + except Exception as e: + logger.error(f"Error enhancing prompt with context: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post("/rag/reindex") +async def reindex_knowledge_base( + organization_id: Optional[str] = Body(None, embed=True), + team_id: Optional[str] = Body(None, embed=True), + agent_id: Optional[str] = Body(None, embed=True), +): + """Reindex all documents in a scope for RAG""" + try: + results = await rag_system.index_all_documents( + organization_id=organization_id, team_id=team_id, agent_id=agent_id + ) + + return { + "scope": { + "organization_id": organization_id, + "team_id": team_id, + "agent_id": agent_id, + }, + "results": results, + "message": f"Indexed {results['indexed']} documents, {results['failed']} failed, {results['skipped']} skipped", + } + + except Exception as e: + logger.error(f"Error reindexing knowledge base: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get("/rag/stats") +async def get_rag_index_stats(): + """Get statistics about the RAG index""" + try: + stats = await rag_system.get_index_stats() + return stats + + except Exception as e: + logger.error(f"Error getting RAG stats: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post("/rag/documents/{doc_id}/reindex") +async def reindex_document( + doc_id: str, + organization_id: Optional[str] = Body(None, embed=True), + team_id: Optional[str] = Body(None, embed=True), + agent_id: Optional[str] = Body(None, embed=True), +): + """Reindex a specific document for RAG""" + try: + # Get document metadata + document = await knowledge_manager.get_document_metadata( + doc_id=doc_id, + organization_id=organization_id, + team_id=team_id, + agent_id=agent_id, + ) + + if not document: + raise HTTPException(status_code=404, detail="Document not found") + + # Reindex the document + success = await rag_system.index_document(document) + + if success: + return { + "document_id": doc_id, + "status": "reindexed", + "message": f"Document '{document.title}' has been reindexed successfully", + } + else: + raise HTTPException(status_code=500, detail="Failed to reindex document") + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error reindexing document {doc_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# ============================================================================ +# Real-time WebSocket Endpoints +# ============================================================================ + + +@app.websocket("/ws/updates") +async def websocket_real_time_updates( + websocket: WebSocket, + organization_id: Optional[str] = None, + team_id: Optional[str] = None, + agent_id: Optional[str] = None, + user_id: Optional[str] = None, + subscriptions: Optional[str] = None, +): + """Main WebSocket endpoint for real-time updates""" + import uuid + + connection_id = str(uuid.uuid4()) + + # Parse subscriptions + subscription_list = [] + if subscriptions: + subscription_list = subscriptions.split(",") + + try: + connection = await websocket_manager.connect( + websocket=websocket, + connection_id=connection_id, + organization_id=organization_id, + team_id=team_id, + agent_id=agent_id, + user_id=user_id, + subscriptions=subscription_list, + ) + + # Keep connection alive and handle pings + while True: + try: + # Wait for ping messages or disconnection + message = await websocket.receive_text() + + # Handle ping/pong + if message == "ping": + await websocket.send_text("pong") + connection.last_ping = datetime.now() + else: + # Parse other messages (subscription updates, etc.) + try: + data = json.loads(message) + if data.get("type") == "subscribe": + # Update subscriptions + new_subs = data.get("subscriptions", []) + connection.scope.subscriptions.clear() + for sub in new_subs: + try: + connection.scope.subscriptions.add(UpdateType(sub)) + except ValueError: + pass + + await connection.send_update( + WebSocketUpdate( + type=UpdateType.SYSTEM_NOTIFICATION, + data={ + "message": "Subscriptions updated", + "subscriptions": list( + connection.scope.subscriptions + ), + }, + ) + ) + except json.JSONDecodeError: + pass + + except WebSocketDisconnect: + break + + except Exception as e: + logger.error(f"WebSocket error for connection {connection_id}: {e}") + finally: + await websocket_manager.disconnect(connection_id) + + +@app.websocket("/ws/agent/{agent_id}/updates") +async def websocket_agent_updates(websocket: WebSocket, agent_id: str): + """WebSocket endpoint for specific agent updates""" + import uuid + + connection_id = f"agent-{agent_id}-{uuid.uuid4()}" + + try: + connection = await websocket_manager.connect( + websocket=websocket, + connection_id=connection_id, + agent_id=agent_id, + subscriptions=[ + UpdateType.AGENT_STATUS.value, + UpdateType.TASK_STATUS.value, + UpdateType.TASK_PROGRESS.value, + UpdateType.CONTAINER_STATUS.value, + UpdateType.CHAT_MESSAGE.value, + UpdateType.CHAT_TYPING.value, + ], + ) + + # Keep connection alive + while True: + try: + message = await websocket.receive_text() + if message == "ping": + await websocket.send_text("pong") + connection.last_ping = datetime.now() + except WebSocketDisconnect: + break + + except Exception as e: + logger.error(f"Agent WebSocket error for {agent_id}: {e}") + finally: + await websocket_manager.disconnect(connection_id) + + +@app.websocket("/ws/agents/{agent_id}/conversations/{conversation_id}") +async def websocket_agent_conversation( + websocket: WebSocket, agent_id: str, conversation_id: str +): + """WebSocket endpoint for real-time agent conversation""" + import uuid + + connection_id = f"conversation-{conversation_id}-{uuid.uuid4()}" + + await websocket.accept() + + try: + # Store connection for broadcasting + active_conversations = getattr(app.state, "active_conversations", {}) + if conversation_id not in active_conversations: + active_conversations[conversation_id] = [] + active_conversations[conversation_id].append(websocket) + app.state.active_conversations = active_conversations + + while True: + try: + # Receive message from client + data = await websocket.receive_json() + + if data.get("type") == "ping": + await websocket.send_json({"type": "pong"}) + elif data.get("type") == "message": + # Handle new message + message_content = data.get("content", "") + if message_content: + # Store message in database + async with get_db_connection() as conn: + message_id = await conn.fetchval( + """ + INSERT INTO agent_conversations (session_id, agent_id, message_type, content) + VALUES ($1, $2, 'user', $3) + RETURNING id + """, + conversation_id, + agent_id, + message_content, + ) + + # Update session activity + await conn.execute( + """ + UPDATE chat_sessions + SET last_activity = CURRENT_TIMESTAMP, message_count = message_count + 1 + WHERE id = $1 + """, + conversation_id, + ) + + # Broadcast to all connected clients for this conversation + message_data = { + "type": "new_message", + "message": { + "id": str(message_id), + "conversation_id": conversation_id, + "role": "user", + "content": message_content, + "timestamp": datetime.now().isoformat(), + "status": "sent", + }, + } + + for conn in active_conversations.get(conversation_id, []): + try: + await conn.send_json(message_data) + except Exception: + # Connection might be closed; skip this subscriber + logger.debug("Skipped broadcast to a closed websocket") + + # TODO: Here we would trigger agent response generation + # For now, send a simple acknowledgment after a delay + await asyncio.sleep(1) + + agent_response = { + "type": "new_message", + "message": { + "id": str(uuid.uuid4()), + "conversation_id": conversation_id, + "role": "agent", + "content": f"I received your message: {message_content}", + "timestamp": datetime.now().isoformat(), + "status": "received", + }, + } + + for conn in active_conversations.get(conversation_id, []): + try: + await conn.send_json(agent_response) + except Exception: + # Connection might be closed; skip this subscriber + logger.debug("Skipped broadcast to a closed websocket") + + # Store agent response in database + async with get_db_connection() as conn: + await conn.execute( + """ + INSERT INTO agent_conversations (session_id, agent_id, message_type, content) + VALUES ($1, $2, 'agent', $3) + """, + conversation_id, + agent_id, + agent_response["message"]["content"], + ) + + except WebSocketDisconnect: + break + except Exception as e: + logger.error(f"Error in conversation WebSocket: {e}") + + except Exception as e: + logger.error(f"Conversation WebSocket error for {conversation_id}: {e}") + finally: + # Clean up connection + if ( + hasattr(app.state, "active_conversations") + and conversation_id in app.state.active_conversations + ): + if websocket in app.state.active_conversations[conversation_id]: + app.state.active_conversations[conversation_id].remove(websocket) + + +@app.websocket("/ws/organization/{organization_id}/updates") +async def websocket_organization_updates(websocket: WebSocket, organization_id: str): + """WebSocket endpoint for organization-wide updates""" + import uuid + + connection_id = f"org-{organization_id}-{uuid.uuid4()}" + + try: + connection = await websocket_manager.connect( + websocket=websocket, + connection_id=connection_id, + organization_id=organization_id, + subscriptions=[ + UpdateType.AGENT_CREATED.value, + UpdateType.AGENT_UPDATED.value, + UpdateType.AGENT_DELETED.value, + UpdateType.KNOWLEDGE_UPDATED.value, + UpdateType.KNOWLEDGE_INDEXED.value, + UpdateType.SYSTEM_NOTIFICATION.value, + ], + ) + + # Keep connection alive + while True: + try: + message = await websocket.receive_text() + if message == "ping": + await websocket.send_text("pong") + connection.last_ping = datetime.now() + except WebSocketDisconnect: + break + + except Exception as e: + logger.error(f"Organization WebSocket error for {organization_id}: {e}") + finally: + await websocket_manager.disconnect(connection_id) + + +# WebSocket Statistics Endpoint +@app.get("/ws/stats") +async def get_websocket_stats(): + """Get WebSocket connection statistics""" + try: + stats = websocket_manager.get_stats() + return stats + except Exception as e: + logger.error(f"Error getting WebSocket stats: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# Manual notification endpoints for testing +@app.post("/ws/test/agent/{agent_id}/status") +async def test_agent_status_notification( + agent_id: str, + status: str = Body(..., embed=True), + message: Optional[str] = Body(None, embed=True), +): + """Test endpoint to send agent status notifications""" + try: + await notify_agent_status_change( + agent_id=agent_id, + status=status, + additional_data={"message": message} if message else None, + ) + return {"status": "notification_sent", "agent_id": agent_id} + except Exception as e: + logger.error(f"Error sending test notification: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# ============================================================================ +# Missing API Endpoints (Goals, Teams, Organizations) +# ============================================================================ + + +@app.get("/teams") +async def get_teams(): + """Get list of teams""" + # Mock data for now + return [ + { + "id": "1", + "name": "Development Team", + "description": "Frontend and backend developers", + "member_count": 5, + "organization_id": "1", + }, + { + "id": "2", + "name": "Executive Team", + "description": "Leadership and strategy", + "member_count": 3, + "organization_id": "1", + }, + ] + + +@app.get("/organizations/{organization_id}/goals") +async def get_organization_goals(organization_id: str): + """Get goals for an organization""" + # Mock data for now + return [ + { + "id": "1", + "title": "Increase Development Velocity", + "description": "Improve team productivity and code quality", + "status": "active", + "progress": 75, + "organization_id": organization_id, + "created_at": "2024-01-15T10:00:00Z", + "due_date": "2024-12-31T23:59:59Z", + }, + { + "id": "2", + "title": "Enhance AI Capabilities", + "description": "Expand AI agent capabilities and intelligence", + "status": "active", + "progress": 50, + "organization_id": organization_id, + "created_at": "2024-02-01T10:00:00Z", + "due_date": "2024-11-30T23:59:59Z", + }, + ] + + +@app.get("/goals/{goal_id}") +async def get_goal_details(goal_id: str): + """Get detailed information about a specific goal""" + # Mock data for now + return { + "id": goal_id, + "title": "Increase Development Velocity", + "description": "Improve team productivity and code quality through better tooling, processes, and automation", + "status": "active", + "progress": 75, + "organization_id": "1", + "team_id": "1", + "created_at": "2024-01-15T10:00:00Z", + "updated_at": "2024-08-06T16:30:00Z", + "due_date": "2024-12-31T23:59:59Z", + "milestones": [ + { + "id": "1", + "title": "Implement CI/CD Pipeline", + "description": "Set up automated testing and deployment", + "status": "completed", + "progress": 100, + "due_date": "2024-03-15T23:59:59Z", + }, + { + "id": "2", + "title": "Enhance Code Review Process", + "description": "Streamline code review workflow with automated tools", + "status": "in_progress", + "progress": 80, + "due_date": "2024-09-30T23:59:59Z", + }, + { + "id": "3", + "title": "Deploy AI-Powered Testing", + "description": "Implement intelligent test generation and execution", + "status": "planned", + "progress": 25, + "due_date": "2024-12-15T23:59:59Z", + }, + ], + "metrics": { + "deployment_frequency": "Daily", + "lead_time": "2.3 days", + "mttr": "45 minutes", + "change_failure_rate": "5%", + }, + "assigned_agents": [ + {"id": "1", "name": "DevOps Agent", "role": "CI/CD Specialist"}, + {"id": "2", "name": "QA Agent", "role": "Test Automation Engineer"}, + ], + } + + +@app.get("/agents/{agent_id}/tasks") +async def get_agent_tasks_list(agent_id: str): + """Get tasks for a specific agent - GET method""" + try: + # Get tasks from task queue + tasks = await app.state.task_queue.get_agent_tasks(agent_id) + return {"agent_id": agent_id, "tasks": tasks} + except Exception as e: + logger.error(f"Error getting tasks for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) diff --git a/services/orchestrator/mcp_integration.py b/services/orchestrator/mcp_integration.py index 34664a4..8393511 100644 --- a/services/orchestrator/mcp_integration.py +++ b/services/orchestrator/mcp_integration.py @@ -1,659 +1,659 @@ -""" -MCP (Model Context Protocol) Integration for FuzeAgent - -Provides MCP server functionality to give Claude SDK sessions access to: -- Organization structure and context -- Team information and agent hierarchy -- Agent capabilities and current status -- Task context and history -- Repository and project information - -This allows agents to have full organizational context when making decisions. -""" - -import asyncio -import json -import logging -import os -import uuid -from dataclasses import asdict, dataclass -from datetime import datetime -from typing import Any, Dict, List, Optional, Union - -from .database import DatabaseManager - -logger = logging.getLogger(__name__) - - -@dataclass -class MCPTool: - """Represents an MCP tool definition""" - - name: str - description: str - input_schema: Dict[str, Any] - - -@dataclass -class MCPResource: - """Represents an MCP resource""" - - uri: str - name: str - description: str - mime_type: str - - -class FuzeAgentMCPServer: - """ - MCP Server for FuzeAgent organizational context. - - Provides tools and resources for Claude SDK sessions to access: - - Organizational structure - - Agent capabilities and status - - Task context and history - - Repository information - """ - - def __init__(self): - self.tools = self._define_tools() - self.resources = self._define_resources() - - def _define_tools(self) -> List[MCPTool]: - """Define available MCP tools""" - return [ - MCPTool( - name="get_organization_structure", - description="Get the complete organizational structure including teams and agents", - input_schema={ - "type": "object", - "properties": { - "organization_id": { - "type": "string", - "description": "Optional organization ID to filter by", - } - }, - }, - ), - MCPTool( - name="get_team_agents", - description="Get all agents in a specific team with their capabilities", - input_schema={ - "type": "object", - "properties": { - "team_id": { - "type": "string", - "description": "Team ID to get agents for", - } - }, - "required": ["team_id"], - }, - ), - MCPTool( - name="get_agent_status", - description="Get current status and capabilities of a specific agent", - input_schema={ - "type": "object", - "properties": { - "agent_id": { - "type": "string", - "description": "Agent ID to get status for", - } - }, - "required": ["agent_id"], - }, - ), - MCPTool( - name="get_task_context", - description="Get comprehensive context for a task including history and related tasks", - input_schema={ - "type": "object", - "properties": { - "task_id": { - "type": "string", - "description": "Task ID to get context for", - }, - "include_history": { - "type": "boolean", - "description": "Whether to include task execution history", - "default": True, - }, - }, - "required": ["task_id"], - }, - ), - MCPTool( - name="get_agent_memory", - description="Get agent memory and previous interactions", - input_schema={ - "type": "object", - "properties": { - "agent_id": { - "type": "string", - "description": "Agent ID to get memory for", - }, - "limit": { - "type": "integer", - "description": "Maximum number of memory items to return", - "default": 10, - }, - "memory_type": { - "type": "string", - "description": "Type of memory to retrieve", - "enum": [ - "interactions", - "code_generations", - "performance_metrics", - ], - "default": "interactions", - }, - }, - "required": ["agent_id"], - }, - ), - MCPTool( - name="search_similar_tasks", - description="Search for similar tasks based on description or requirements", - input_schema={ - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "Search query for similar tasks", - }, - "agent_type": { - "type": "string", - "description": "Optional agent type to filter results", - }, - "limit": { - "type": "integer", - "description": "Maximum number of results", - "default": 5, - }, - }, - "required": ["query"], - }, - ), - MCPTool( - name="get_repository_context", - description="Get repository context and recent changes", - input_schema={ - "type": "object", - "properties": { - "repository_url": { - "type": "string", - "description": "Repository URL to get context for", - }, - "branch": { - "type": "string", - "description": "Optional branch name", - "default": "main", - }, - }, - "required": ["repository_url"], - }, - ), - MCPTool( - name="get_agent_recommendations", - description="Get agent recommendations for a specific task type", - input_schema={ - "type": "object", - "properties": { - "task_description": { - "type": "string", - "description": "Description of the task", - }, - "required_skills": { - "type": "array", - "items": {"type": "string"}, - "description": "List of required skills", - }, - "exclude_busy": { - "type": "boolean", - "description": "Whether to exclude currently busy agents", - "default": True, - }, - }, - "required": ["task_description"], - }, - ), - ] - - def _define_resources(self) -> List[MCPResource]: - """Define available MCP resources""" - return [ - MCPResource( - uri="fuzeagent://organizations", - name="Organizations", - description="Complete organizational structure and hierarchy", - mime_type="application/json", - ), - MCPResource( - uri="fuzeagent://agent-templates", - name="Agent Templates", - description="Available agent templates and their capabilities", - mime_type="application/json", - ), - MCPResource( - uri="fuzeagent://system-status", - name="System Status", - description="Current system status and health metrics", - mime_type="application/json", - ), - ] - - async def handle_tool_call( - self, tool_name: str, arguments: Dict[str, Any] - ) -> Dict[str, Any]: - """Handle MCP tool calls""" - try: - if tool_name == "get_organization_structure": - return await self._get_organization_structure(arguments) - elif tool_name == "get_team_agents": - return await self._get_team_agents(arguments) - elif tool_name == "get_agent_status": - return await self._get_agent_status(arguments) - elif tool_name == "get_task_context": - return await self._get_task_context(arguments) - elif tool_name == "get_agent_memory": - return await self._get_agent_memory(arguments) - elif tool_name == "search_similar_tasks": - return await self._search_similar_tasks(arguments) - elif tool_name == "get_repository_context": - return await self._get_repository_context(arguments) - elif tool_name == "get_agent_recommendations": - return await self._get_agent_recommendations(arguments) - else: - return {"error": f"Unknown tool: {tool_name}"} - - except Exception as e: - logger.error(f"Error in MCP tool call {tool_name}: {e}") - return {"error": str(e)} - - async def handle_resource_request(self, uri: str) -> Dict[str, Any]: - """Handle MCP resource requests""" - try: - if uri == "fuzeagent://organizations": - return await self._get_organizations_resource() - elif uri == "fuzeagent://agent-templates": - return await self._get_agent_templates_resource() - elif uri == "fuzeagent://system-status": - return await self._get_system_status_resource() - else: - return {"error": f"Unknown resource: {uri}"} - - except Exception as e: - logger.error(f"Error in MCP resource request {uri}: {e}") - return {"error": str(e)} - - # Tool implementations - - async def _get_organization_structure(self, args: Dict[str, Any]) -> Dict[str, Any]: - """Get organizational structure""" - organization_id = args.get("organization_id") - - # This would integrate with the MCP FuzeAgent server - # For now, return mock structure - return { - "organizations": [ - { - "id": "fuzeagent-org", - "name": "FuzeAgent Organization", - "teams": [ - { - "id": "dev-team-1", - "name": "Development Team Alpha", - "agents": [ - { - "id": "frontend-dev-1", - "name": "React Developer 1", - "type": "frontend_developer", - "status": "available", - "skills": ["react", "typescript", "css"], - }, - { - "id": "backend-dev-1", - "name": "Python Developer 1", - "type": "backend_developer", - "status": "busy", - "skills": ["python", "fastapi", "postgresql"], - }, - ], - } - ], - } - ] - } - - async def _get_team_agents(self, args: Dict[str, Any]) -> Dict[str, Any]: - """Get agents in a team""" - team_id = args["team_id"] - - # Get agents from database - agents = await DatabaseManager.get_agents_by_team(team_id) - - return { - "team_id": team_id, - "agents": [ - { - "id": agent["id"], - "name": agent["name"], - "role": agent["role"], - "type": agent["type"], - "status": agent["status"], - "capabilities": agent.get("config", {}).get("tools", []), - "current_task": agent.get("current_task_id"), - "created_at": ( - agent["created_at"].isoformat() if agent["created_at"] else None - ), - } - for agent in agents - ], - } - - async def _get_agent_status(self, args: Dict[str, Any]) -> Dict[str, Any]: - """Get agent status""" - agent_id = args["agent_id"] - - agent = await DatabaseManager.get_agent(agent_id) - if not agent: - return {"error": f"Agent {agent_id} not found"} - - # Get current tasks - tasks = await DatabaseManager.get_agent_tasks(agent_id, limit=5) - - return { - "agent_id": agent_id, - "name": agent["name"], - "role": agent["role"], - "type": agent["type"], - "status": agent["status"], - "capabilities": agent.get("config", {}).get("tools", []), - "model": agent.get("config", {}).get("model", "claude-sonnet-4-20250514"), - "current_tasks": [ - { - "id": task["id"], - "title": task["title"], - "status": task["status"], - "created_at": ( - task["created_at"].isoformat() if task["created_at"] else None - ), - } - for task in tasks - if task["status"] in ["pending", "executing"] - ], - "recent_tasks": [ - { - "id": task["id"], - "title": task["title"], - "status": task["status"], - "completed_at": ( - task["updated_at"].isoformat() if task["updated_at"] else None - ), - } - for task in tasks - if task["status"] in ["completed", "failed"] - ], - } - - async def _get_task_context(self, args: Dict[str, Any]) -> Dict[str, Any]: - """Get task context""" - task_id = args["task_id"] - include_history = args.get("include_history", True) - - # Get task data - task = await DatabaseManager.get_task(task_id) - if not task: - return {"error": f"Task {task_id} not found"} - - # Get agent data - agent = ( - await DatabaseManager.get_agent(task["assigned_to"]) - if task["assigned_to"] - else None - ) - - context = { - "task_id": task_id, - "title": task["title"], - "description": task["description"], - "status": task["status"], - "priority": task.get("priority", "medium"), - "created_at": ( - task["created_at"].isoformat() if task["created_at"] else None - ), - "assigned_agent": ( - { - "id": agent["id"], - "name": agent["name"], - "role": agent["role"], - "type": agent["type"], - } - if agent - else None - ), - } - - if include_history: - # Get task iterations - iterations = await DatabaseManager.get_task_iterations(task_id) - context["execution_history"] = [ - { - "iteration": iter["iteration_number"], - "step": iter["step"], - "started_at": ( - iter["started_at"].isoformat() if iter["started_at"] else None - ), - "completed_at": ( - iter["completed_at"].isoformat() - if iter["completed_at"] - else None - ), - "success": iter["success"], - "human_question": iter["human_question"], - "human_response": iter["human_response"], - } - for iter in iterations - ] - - return context - - async def _get_agent_memory(self, args: Dict[str, Any]) -> Dict[str, Any]: - """Get agent memory""" - agent_id = args["agent_id"] - limit = args.get("limit", 10) - memory_type = args.get("memory_type", "interactions") - - # This would integrate with conversation_manager - # For now return mock data - return { - "agent_id": agent_id, - "memory_type": memory_type, - "items": [ - { - "id": f"memory-{i}", - "type": memory_type, - "content": f"Sample {memory_type} {i}", - "timestamp": datetime.now().isoformat(), - "metadata": {}, - } - for i in range(min(limit, 5)) - ], - } - - async def _search_similar_tasks(self, args: Dict[str, Any]) -> Dict[str, Any]: - """Search for similar tasks""" - query = args["query"] - agent_type = args.get("agent_type") - limit = args.get("limit", 5) - - # This would use vector search in production - # For now return mock results - return { - "query": query, - "results": [ - { - "task_id": f"task-{i}", - "title": f"Similar task {i}", - "description": f"Task similar to '{query}'", - "similarity_score": 0.8 - (i * 0.1), - "agent_type": agent_type or "developer", - "status": "completed", - "completion_time_minutes": 120 + (i * 30), - } - for i in range(min(limit, 3)) - ], - } - - async def _get_repository_context(self, args: Dict[str, Any]) -> Dict[str, Any]: - """Get repository context""" - repository_url = args["repository_url"] - branch = args.get("branch", "main") - - return { - "repository_url": repository_url, - "branch": branch, - "recent_commits": [ - { - "hash": "abc123", - "message": "Recent commit message", - "author": "developer@example.com", - "timestamp": datetime.now().isoformat(), - } - ], - "active_branches": [branch, "develop", "feature/new-feature"], - "technologies": ["python", "fastapi", "react", "typescript"], - "structure": { - "backend": "services/orchestrator/", - "frontend": "services/ui-react/", - "containers": "containers/", - "docs": "docs/", - }, - } - - async def _get_agent_recommendations(self, args: Dict[str, Any]) -> Dict[str, Any]: - """Get agent recommendations for a task""" - task_description = args["task_description"] - required_skills = args.get("required_skills", []) - exclude_busy = args.get("exclude_busy", True) - - # This would use ML/AI to match agents to tasks - # For now return mock recommendations - return { - "task_description": task_description, - "required_skills": required_skills, - "recommendations": [ - { - "agent_id": "frontend-dev-1", - "name": "React Developer 1", - "match_score": 0.95, - "matching_skills": ["react", "typescript"], - "availability": "available", - "estimated_completion_time": "4-6 hours", - }, - { - "agent_id": "fullstack-dev-1", - "name": "Full Stack Developer 1", - "match_score": 0.85, - "matching_skills": ["react", "python"], - "availability": "busy_until_2pm", - "estimated_completion_time": "6-8 hours", - }, - ], - } - - # Resource implementations - - async def _get_organizations_resource(self) -> Dict[str, Any]: - """Get organizations resource""" - organizations = await DatabaseManager.get_organizations() - return {"content": organizations, "mime_type": "application/json"} - - async def _get_agent_templates_resource(self) -> Dict[str, Any]: - """Get agent templates resource""" - templates = await DatabaseManager.get_agent_templates() - return {"content": templates, "mime_type": "application/json"} - - async def _get_system_status_resource(self) -> Dict[str, Any]: - """Get system status resource""" - return { - "content": { - "status": "healthy", - "timestamp": datetime.now().isoformat(), - "active_agents": 5, - "running_tasks": 3, - "system_load": 0.45, - "memory_usage": 0.67, - }, - "mime_type": "application/json", - } - - -# MCP Server Integration with Claude SDK -class MCPClaudeIntegration: - """Integrates MCP server with Claude SDK sessions""" - - def __init__(self, mcp_server: FuzeAgentMCPServer): - self.mcp_server = mcp_server - - async def setup_claude_session_mcp( - self, session_id: str, agent_id: str, task_id: str - ) -> Dict[str, str]: - """Set up MCP tools for a Claude SDK session""" - - # Generate MCP server configuration for the session - mcp_config = { - "server_name": f"fuzeagent-{session_id}", - "server_command": ["python", "-m", "services.orchestrator.mcp_integration"], - "server_args": [ - "--session-id", - session_id, - "--agent-id", - agent_id, - "--task-id", - task_id, - ], - "environment": { - "FUZEAGENT_SESSION_ID": session_id, - "FUZEAGENT_AGENT_ID": agent_id, - "FUZEAGENT_TASK_ID": task_id, - }, - } - - # In production, this would configure Claude SDK to use this MCP server - # For now, return the configuration - return mcp_config - - async def get_session_context( - self, session_id: str, agent_id: str, task_id: str - ) -> Dict[str, Any]: - """Get comprehensive context for a Claude SDK session""" - - # Get agent context - agent_context = await self.mcp_server.handle_tool_call( - "get_agent_status", {"agent_id": agent_id} - ) - - # Get task context - task_context = await self.mcp_server.handle_tool_call( - "get_task_context", {"task_id": task_id} - ) - - # Get team context - if agent_context.get("team_id"): - team_context = await self.mcp_server.handle_tool_call( - "get_team_agents", {"team_id": agent_context["team_id"]} - ) - else: - team_context = {"agents": []} - - return { - "session_id": session_id, - "agent_context": agent_context, - "task_context": task_context, - "team_context": team_context, - "available_tools": [tool.name for tool in self.mcp_server.tools], - "available_resources": [ - resource.uri for resource in self.mcp_server.resources - ], - } +""" +MCP (Model Context Protocol) Integration for FuzeAgent + +Provides MCP server functionality to give Claude SDK sessions access to: +- Organization structure and context +- Team information and agent hierarchy +- Agent capabilities and current status +- Task context and history +- Repository and project information + +This allows agents to have full organizational context when making decisions. +""" + +import asyncio +import json +import logging +import os +import uuid +from dataclasses import asdict, dataclass +from datetime import datetime +from typing import Any, Dict, List, Optional, Union + +from .database import DatabaseManager + +logger = logging.getLogger(__name__) + + +@dataclass +class MCPTool: + """Represents an MCP tool definition""" + + name: str + description: str + input_schema: Dict[str, Any] + + +@dataclass +class MCPResource: + """Represents an MCP resource""" + + uri: str + name: str + description: str + mime_type: str + + +class FuzeAgentMCPServer: + """ + MCP Server for FuzeAgent organizational context. + + Provides tools and resources for Claude SDK sessions to access: + - Organizational structure + - Agent capabilities and status + - Task context and history + - Repository information + """ + + def __init__(self): + self.tools = self._define_tools() + self.resources = self._define_resources() + + def _define_tools(self) -> List[MCPTool]: + """Define available MCP tools""" + return [ + MCPTool( + name="get_organization_structure", + description="Get the complete organizational structure including teams and agents", + input_schema={ + "type": "object", + "properties": { + "organization_id": { + "type": "string", + "description": "Optional organization ID to filter by", + } + }, + }, + ), + MCPTool( + name="get_team_agents", + description="Get all agents in a specific team with their capabilities", + input_schema={ + "type": "object", + "properties": { + "team_id": { + "type": "string", + "description": "Team ID to get agents for", + } + }, + "required": ["team_id"], + }, + ), + MCPTool( + name="get_agent_status", + description="Get current status and capabilities of a specific agent", + input_schema={ + "type": "object", + "properties": { + "agent_id": { + "type": "string", + "description": "Agent ID to get status for", + } + }, + "required": ["agent_id"], + }, + ), + MCPTool( + name="get_task_context", + description="Get comprehensive context for a task including history and related tasks", + input_schema={ + "type": "object", + "properties": { + "task_id": { + "type": "string", + "description": "Task ID to get context for", + }, + "include_history": { + "type": "boolean", + "description": "Whether to include task execution history", + "default": True, + }, + }, + "required": ["task_id"], + }, + ), + MCPTool( + name="get_agent_memory", + description="Get agent memory and previous interactions", + input_schema={ + "type": "object", + "properties": { + "agent_id": { + "type": "string", + "description": "Agent ID to get memory for", + }, + "limit": { + "type": "integer", + "description": "Maximum number of memory items to return", + "default": 10, + }, + "memory_type": { + "type": "string", + "description": "Type of memory to retrieve", + "enum": [ + "interactions", + "code_generations", + "performance_metrics", + ], + "default": "interactions", + }, + }, + "required": ["agent_id"], + }, + ), + MCPTool( + name="search_similar_tasks", + description="Search for similar tasks based on description or requirements", + input_schema={ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Search query for similar tasks", + }, + "agent_type": { + "type": "string", + "description": "Optional agent type to filter results", + }, + "limit": { + "type": "integer", + "description": "Maximum number of results", + "default": 5, + }, + }, + "required": ["query"], + }, + ), + MCPTool( + name="get_repository_context", + description="Get repository context and recent changes", + input_schema={ + "type": "object", + "properties": { + "repository_url": { + "type": "string", + "description": "Repository URL to get context for", + }, + "branch": { + "type": "string", + "description": "Optional branch name", + "default": "main", + }, + }, + "required": ["repository_url"], + }, + ), + MCPTool( + name="get_agent_recommendations", + description="Get agent recommendations for a specific task type", + input_schema={ + "type": "object", + "properties": { + "task_description": { + "type": "string", + "description": "Description of the task", + }, + "required_skills": { + "type": "array", + "items": {"type": "string"}, + "description": "List of required skills", + }, + "exclude_busy": { + "type": "boolean", + "description": "Whether to exclude currently busy agents", + "default": True, + }, + }, + "required": ["task_description"], + }, + ), + ] + + def _define_resources(self) -> List[MCPResource]: + """Define available MCP resources""" + return [ + MCPResource( + uri="fuzeagent://organizations", + name="Organizations", + description="Complete organizational structure and hierarchy", + mime_type="application/json", + ), + MCPResource( + uri="fuzeagent://agent-templates", + name="Agent Templates", + description="Available agent templates and their capabilities", + mime_type="application/json", + ), + MCPResource( + uri="fuzeagent://system-status", + name="System Status", + description="Current system status and health metrics", + mime_type="application/json", + ), + ] + + async def handle_tool_call( + self, tool_name: str, arguments: Dict[str, Any] + ) -> Dict[str, Any]: + """Handle MCP tool calls""" + try: + if tool_name == "get_organization_structure": + return await self._get_organization_structure(arguments) + elif tool_name == "get_team_agents": + return await self._get_team_agents(arguments) + elif tool_name == "get_agent_status": + return await self._get_agent_status(arguments) + elif tool_name == "get_task_context": + return await self._get_task_context(arguments) + elif tool_name == "get_agent_memory": + return await self._get_agent_memory(arguments) + elif tool_name == "search_similar_tasks": + return await self._search_similar_tasks(arguments) + elif tool_name == "get_repository_context": + return await self._get_repository_context(arguments) + elif tool_name == "get_agent_recommendations": + return await self._get_agent_recommendations(arguments) + else: + return {"error": f"Unknown tool: {tool_name}"} + + except Exception as e: + logger.error(f"Error in MCP tool call {tool_name}: {e}") + return {"error": str(e)} + + async def handle_resource_request(self, uri: str) -> Dict[str, Any]: + """Handle MCP resource requests""" + try: + if uri == "fuzeagent://organizations": + return await self._get_organizations_resource() + elif uri == "fuzeagent://agent-templates": + return await self._get_agent_templates_resource() + elif uri == "fuzeagent://system-status": + return await self._get_system_status_resource() + else: + return {"error": f"Unknown resource: {uri}"} + + except Exception as e: + logger.error(f"Error in MCP resource request {uri}: {e}") + return {"error": str(e)} + + # Tool implementations + + async def _get_organization_structure(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Get organizational structure""" + organization_id = args.get("organization_id") + + # This would integrate with the MCP FuzeAgent server + # For now, return mock structure + return { + "organizations": [ + { + "id": "fuzeagent-org", + "name": "FuzeAgent Organization", + "teams": [ + { + "id": "dev-team-1", + "name": "Development Team Alpha", + "agents": [ + { + "id": "frontend-dev-1", + "name": "React Developer 1", + "type": "frontend_developer", + "status": "available", + "skills": ["react", "typescript", "css"], + }, + { + "id": "backend-dev-1", + "name": "Python Developer 1", + "type": "backend_developer", + "status": "busy", + "skills": ["python", "fastapi", "postgresql"], + }, + ], + } + ], + } + ] + } + + async def _get_team_agents(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Get agents in a team""" + team_id = args["team_id"] + + # Get agents from database + agents = await DatabaseManager.get_agents_by_team(team_id) + + return { + "team_id": team_id, + "agents": [ + { + "id": agent["id"], + "name": agent["name"], + "role": agent["role"], + "type": agent["type"], + "status": agent["status"], + "capabilities": agent.get("config", {}).get("tools", []), + "current_task": agent.get("current_task_id"), + "created_at": ( + agent["created_at"].isoformat() if agent["created_at"] else None + ), + } + for agent in agents + ], + } + + async def _get_agent_status(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Get agent status""" + agent_id = args["agent_id"] + + agent = await DatabaseManager.get_agent(agent_id) + if not agent: + return {"error": f"Agent {agent_id} not found"} + + # Get current tasks + tasks = await DatabaseManager.get_agent_tasks(agent_id, limit=5) + + return { + "agent_id": agent_id, + "name": agent["name"], + "role": agent["role"], + "type": agent["type"], + "status": agent["status"], + "capabilities": agent.get("config", {}).get("tools", []), + "model": agent.get("config", {}).get("model", "claude-sonnet-4-20250514"), + "current_tasks": [ + { + "id": task["id"], + "title": task["title"], + "status": task["status"], + "created_at": ( + task["created_at"].isoformat() if task["created_at"] else None + ), + } + for task in tasks + if task["status"] in ["pending", "executing"] + ], + "recent_tasks": [ + { + "id": task["id"], + "title": task["title"], + "status": task["status"], + "completed_at": ( + task["updated_at"].isoformat() if task["updated_at"] else None + ), + } + for task in tasks + if task["status"] in ["completed", "failed"] + ], + } + + async def _get_task_context(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Get task context""" + task_id = args["task_id"] + include_history = args.get("include_history", True) + + # Get task data + task = await DatabaseManager.get_task(task_id) + if not task: + return {"error": f"Task {task_id} not found"} + + # Get agent data + agent = ( + await DatabaseManager.get_agent(task["assigned_to"]) + if task["assigned_to"] + else None + ) + + context = { + "task_id": task_id, + "title": task["title"], + "description": task["description"], + "status": task["status"], + "priority": task.get("priority", "medium"), + "created_at": ( + task["created_at"].isoformat() if task["created_at"] else None + ), + "assigned_agent": ( + { + "id": agent["id"], + "name": agent["name"], + "role": agent["role"], + "type": agent["type"], + } + if agent + else None + ), + } + + if include_history: + # Get task iterations + iterations = await DatabaseManager.get_task_iterations(task_id) + context["execution_history"] = [ + { + "iteration": iter["iteration_number"], + "step": iter["step"], + "started_at": ( + iter["started_at"].isoformat() if iter["started_at"] else None + ), + "completed_at": ( + iter["completed_at"].isoformat() + if iter["completed_at"] + else None + ), + "success": iter["success"], + "human_question": iter["human_question"], + "human_response": iter["human_response"], + } + for iter in iterations + ] + + return context + + async def _get_agent_memory(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Get agent memory""" + agent_id = args["agent_id"] + limit = args.get("limit", 10) + memory_type = args.get("memory_type", "interactions") + + # This would integrate with conversation_manager + # For now return mock data + return { + "agent_id": agent_id, + "memory_type": memory_type, + "items": [ + { + "id": f"memory-{i}", + "type": memory_type, + "content": f"Sample {memory_type} {i}", + "timestamp": datetime.now().isoformat(), + "metadata": {}, + } + for i in range(min(limit, 5)) + ], + } + + async def _search_similar_tasks(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Search for similar tasks""" + query = args["query"] + agent_type = args.get("agent_type") + limit = args.get("limit", 5) + + # This would use vector search in production + # For now return mock results + return { + "query": query, + "results": [ + { + "task_id": f"task-{i}", + "title": f"Similar task {i}", + "description": f"Task similar to '{query}'", + "similarity_score": 0.8 - (i * 0.1), + "agent_type": agent_type or "developer", + "status": "completed", + "completion_time_minutes": 120 + (i * 30), + } + for i in range(min(limit, 3)) + ], + } + + async def _get_repository_context(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Get repository context""" + repository_url = args["repository_url"] + branch = args.get("branch", "main") + + return { + "repository_url": repository_url, + "branch": branch, + "recent_commits": [ + { + "hash": "abc123", + "message": "Recent commit message", + "author": "developer@example.com", + "timestamp": datetime.now().isoformat(), + } + ], + "active_branches": [branch, "develop", "feature/new-feature"], + "technologies": ["python", "fastapi", "react", "typescript"], + "structure": { + "backend": "services/orchestrator/", + "frontend": "services/ui-react/", + "containers": "containers/", + "docs": "docs/", + }, + } + + async def _get_agent_recommendations(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Get agent recommendations for a task""" + task_description = args["task_description"] + required_skills = args.get("required_skills", []) + exclude_busy = args.get("exclude_busy", True) + + # This would use ML/AI to match agents to tasks + # For now return mock recommendations + return { + "task_description": task_description, + "required_skills": required_skills, + "recommendations": [ + { + "agent_id": "frontend-dev-1", + "name": "React Developer 1", + "match_score": 0.95, + "matching_skills": ["react", "typescript"], + "availability": "available", + "estimated_completion_time": "4-6 hours", + }, + { + "agent_id": "fullstack-dev-1", + "name": "Full Stack Developer 1", + "match_score": 0.85, + "matching_skills": ["react", "python"], + "availability": "busy_until_2pm", + "estimated_completion_time": "6-8 hours", + }, + ], + } + + # Resource implementations + + async def _get_organizations_resource(self) -> Dict[str, Any]: + """Get organizations resource""" + organizations = await DatabaseManager.get_organizations() + return {"content": organizations, "mime_type": "application/json"} + + async def _get_agent_templates_resource(self) -> Dict[str, Any]: + """Get agent templates resource""" + templates = await DatabaseManager.get_agent_templates() + return {"content": templates, "mime_type": "application/json"} + + async def _get_system_status_resource(self) -> Dict[str, Any]: + """Get system status resource""" + return { + "content": { + "status": "healthy", + "timestamp": datetime.now().isoformat(), + "active_agents": 5, + "running_tasks": 3, + "system_load": 0.45, + "memory_usage": 0.67, + }, + "mime_type": "application/json", + } + + +# MCP Server Integration with Claude SDK +class MCPClaudeIntegration: + """Integrates MCP server with Claude SDK sessions""" + + def __init__(self, mcp_server: FuzeAgentMCPServer): + self.mcp_server = mcp_server + + async def setup_claude_session_mcp( + self, session_id: str, agent_id: str, task_id: str + ) -> Dict[str, str]: + """Set up MCP tools for a Claude SDK session""" + + # Generate MCP server configuration for the session + mcp_config = { + "server_name": f"fuzeagent-{session_id}", + "server_command": ["python", "-m", "services.orchestrator.mcp_integration"], + "server_args": [ + "--session-id", + session_id, + "--agent-id", + agent_id, + "--task-id", + task_id, + ], + "environment": { + "FUZEAGENT_SESSION_ID": session_id, + "FUZEAGENT_AGENT_ID": agent_id, + "FUZEAGENT_TASK_ID": task_id, + }, + } + + # In production, this would configure Claude SDK to use this MCP server + # For now, return the configuration + return mcp_config + + async def get_session_context( + self, session_id: str, agent_id: str, task_id: str + ) -> Dict[str, Any]: + """Get comprehensive context for a Claude SDK session""" + + # Get agent context + agent_context = await self.mcp_server.handle_tool_call( + "get_agent_status", {"agent_id": agent_id} + ) + + # Get task context + task_context = await self.mcp_server.handle_tool_call( + "get_task_context", {"task_id": task_id} + ) + + # Get team context + if agent_context.get("team_id"): + team_context = await self.mcp_server.handle_tool_call( + "get_team_agents", {"team_id": agent_context["team_id"]} + ) + else: + team_context = {"agents": []} + + return { + "session_id": session_id, + "agent_context": agent_context, + "task_context": task_context, + "team_context": team_context, + "available_tools": [tool.name for tool in self.mcp_server.tools], + "available_resources": [ + resource.uri for resource in self.mcp_server.resources + ], + } diff --git a/services/orchestrator/model_configuration.py b/services/orchestrator/model_configuration.py index de5cf16..e0834b1 100644 --- a/services/orchestrator/model_configuration.py +++ b/services/orchestrator/model_configuration.py @@ -1,558 +1,558 @@ -""" -Model Configuration and API Key Management for FuzeAgent - -Manages model configurations and API keys for different AI providers at the -organization level, with secure storage and agent-specific model selection. - -Supports: -- Multiple AI model providers (Anthropic, OpenAI, Google, etc.) -- Organization-level API key management -- Agent-specific model configurations -- Secure credential storage and access -- Model capability matching -- Cost optimization -""" - -import base64 -import json -import logging -import os -import tempfile -from dataclasses import dataclass, field -from datetime import datetime -from enum import Enum -from typing import Any, Dict, List, Optional, Union - -from cryptography.fernet import Fernet - -from .database import DatabaseManager - -logger = logging.getLogger(__name__) - - -class ModelProvider(str, Enum): - ANTHROPIC = "anthropic" - OPENAI = "openai" - GOOGLE = "google" - AZURE_OPENAI = "azure_openai" - COHERE = "cohere" - HUGGINGFACE = "huggingface" - OLLAMA = "ollama" - CUSTOM = "custom" - - -class ModelCapability(str, Enum): - TEXT_GENERATION = "text_generation" - CODE_GENERATION = "code_generation" - REASONING = "reasoning" - ANALYSIS = "analysis" - CONVERSATION = "conversation" - FUNCTION_CALLING = "function_calling" - VISION = "vision" - EMBEDDINGS = "embeddings" - - -@dataclass -class ModelSpec: - """Specification for an AI model""" - - model_id: str - provider: ModelProvider - name: str - description: str - capabilities: List[ModelCapability] - context_window: int - max_output_tokens: int - cost_per_input_token: float # USD per 1K tokens - cost_per_output_token: float # USD per 1K tokens - supports_streaming: bool = True - supports_function_calling: bool = False - supports_vision: bool = False - supports_json_mode: bool = False - temperature_range: tuple = (0.0, 2.0) - recommended_use_cases: List[str] = field(default_factory=list) - - -@dataclass -class ProviderCredentials: - """Secure storage for provider API credentials""" - - provider: ModelProvider - encrypted_api_key: str - endpoint_url: Optional[str] = None - additional_config: Dict[str, Any] = field(default_factory=dict) - created_at: datetime = field(default_factory=datetime.now) - last_used: Optional[datetime] = None - is_active: bool = True - - -@dataclass -class AgentModelConfig: - """Model configuration for an agent""" - - agent_id: str - primary_model: str # model_id - fallback_models: List[str] = field(default_factory=list) - temperature: float = 0.7 - max_tokens: int = 4096 - top_p: float = 1.0 - frequency_penalty: float = 0.0 - presence_penalty: float = 0.0 - custom_instructions: str = "" - use_function_calling: bool = True - streaming_enabled: bool = True - cost_limit_per_task: Optional[float] = None # USD - created_at: datetime = field(default_factory=datetime.now) - updated_at: datetime = field(default_factory=datetime.now) - - -class ModelConfigurationManager: - """ - Manages AI model configurations and provider credentials. - - Features: - - Secure API key storage with encryption - - Organization-level credential management - - Agent-specific model configurations - - Model capability matching - - Cost tracking and limits - - Automatic fallback handling - """ - - def __init__(self): - self.encryption_key = self._get_or_create_encryption_key() - self.fernet = Fernet(self.encryption_key) - self.available_models = self._initialize_model_specs() - - def _get_or_create_encryption_key(self) -> bytes: - """Get or create encryption key for API credentials""" - # Avoid a hard-coded world-readable /tmp path; allow override and fall - # back to the platform temp dir (still /tmp inside the Linux container). - key_dir = os.getenv("FUZEAGENT_KEY_DIR", tempfile.gettempdir()) - key_file = os.path.join(key_dir, "fuzeagent_encryption.key") - - if os.path.exists(key_file): - with open(key_file, "rb") as f: - return f.read() - else: - key = Fernet.generate_key() - with open(key_file, "wb") as f: - f.write(key) - return key - - def _initialize_model_specs(self) -> Dict[str, ModelSpec]: - """Initialize available model specifications""" - models = {} - - # Anthropic Claude models - models["claude-3-5-sonnet-20241022"] = ModelSpec( - model_id="claude-3-5-sonnet-20241022", - provider=ModelProvider.ANTHROPIC, - name="Claude 3.5 Sonnet", - description="Most intelligent model for complex reasoning and coding", - capabilities=[ - ModelCapability.TEXT_GENERATION, - ModelCapability.CODE_GENERATION, - ModelCapability.REASONING, - ModelCapability.ANALYSIS, - ModelCapability.CONVERSATION, - ModelCapability.FUNCTION_CALLING, - ModelCapability.VISION, - ], - context_window=200000, - max_output_tokens=8192, - cost_per_input_token=0.003, - cost_per_output_token=0.015, - supports_function_calling=True, - supports_vision=True, - supports_json_mode=True, - recommended_use_cases=[ - "complex coding", - "reasoning", - "analysis", - "research", - ], - ) - - models["claude-3-haiku-20240307"] = ModelSpec( - model_id="claude-3-haiku-20240307", - provider=ModelProvider.ANTHROPIC, - name="Claude 3 Haiku", - description="Fastest and most cost-effective model for simple tasks", - capabilities=[ - ModelCapability.TEXT_GENERATION, - ModelCapability.CODE_GENERATION, - ModelCapability.CONVERSATION, - ], - context_window=200000, - max_output_tokens=4096, - cost_per_input_token=0.00025, - cost_per_output_token=0.00125, - supports_function_calling=True, - supports_vision=True, - recommended_use_cases=[ - "simple tasks", - "quick responses", - "cost optimization", - ], - ) - - # OpenAI GPT models - models["gpt-4o"] = ModelSpec( - model_id="gpt-4o", - provider=ModelProvider.OPENAI, - name="GPT-4 Omni", - description="OpenAI's most capable multimodal model", - capabilities=[ - ModelCapability.TEXT_GENERATION, - ModelCapability.CODE_GENERATION, - ModelCapability.REASONING, - ModelCapability.ANALYSIS, - ModelCapability.CONVERSATION, - ModelCapability.FUNCTION_CALLING, - ModelCapability.VISION, - ], - context_window=128000, - max_output_tokens=4096, - cost_per_input_token=0.005, - cost_per_output_token=0.015, - supports_function_calling=True, - supports_vision=True, - supports_json_mode=True, - recommended_use_cases=[ - "multimodal tasks", - "function calling", - "complex reasoning", - ], - ) - - models["gpt-4o-mini"] = ModelSpec( - model_id="gpt-4o-mini", - provider=ModelProvider.OPENAI, - name="GPT-4 Omni Mini", - description="Cost-effective model for simpler tasks", - capabilities=[ - ModelCapability.TEXT_GENERATION, - ModelCapability.CODE_GENERATION, - ModelCapability.CONVERSATION, - ModelCapability.FUNCTION_CALLING, - ], - context_window=128000, - max_output_tokens=16384, - cost_per_input_token=0.00015, - cost_per_output_token=0.0006, - supports_function_calling=True, - supports_json_mode=True, - recommended_use_cases=["cost optimization", "simple tasks", "high volume"], - ) - - # Google Gemini models - models["gemini-1.5-pro"] = ModelSpec( - model_id="gemini-1.5-pro", - provider=ModelProvider.GOOGLE, - name="Gemini 1.5 Pro", - description="Google's most capable model with long context", - capabilities=[ - ModelCapability.TEXT_GENERATION, - ModelCapability.CODE_GENERATION, - ModelCapability.REASONING, - ModelCapability.ANALYSIS, - ModelCapability.CONVERSATION, - ModelCapability.FUNCTION_CALLING, - ModelCapability.VISION, - ], - context_window=2000000, # 2M tokens - max_output_tokens=8192, - cost_per_input_token=0.00125, - cost_per_output_token=0.005, - supports_function_calling=True, - supports_vision=True, - recommended_use_cases=[ - "long context", - "document analysis", - "multimodal tasks", - ], - ) - - return models - - async def store_provider_credentials( - self, - organization_id: str, - provider: ModelProvider, - api_key: str, - endpoint_url: Optional[str] = None, - additional_config: Optional[Dict[str, Any]] = None, - ) -> bool: - """Store encrypted API credentials for a provider""" - try: - # Encrypt the API key - encrypted_key = self.fernet.encrypt(api_key.encode()).decode() - - credentials = ProviderCredentials( - provider=provider, - encrypted_api_key=encrypted_key, - endpoint_url=endpoint_url, - additional_config=additional_config or {}, - ) - - # Store in database - await DatabaseManager.store_provider_credentials( - organization_id, credentials.__dict__ - ) - - logger.info( - f"Stored credentials for {provider} in organization {organization_id}" - ) - return True - - except Exception as e: - logger.error(f"Error storing provider credentials: {e}") - return False - - async def get_provider_credentials( - self, organization_id: str, provider: ModelProvider - ) -> Optional[str]: - """Get decrypted API key for a provider""" - try: - credentials_data = await DatabaseManager.get_provider_credentials( - organization_id, provider.value - ) - - if not credentials_data: - return None - - # Decrypt the API key - encrypted_key = credentials_data["encrypted_api_key"] - decrypted_key = self.fernet.decrypt(encrypted_key.encode()).decode() - - # Update last used timestamp - await DatabaseManager.update_credentials_last_used( - organization_id, provider.value - ) - - return decrypted_key - - except Exception as e: - logger.error(f"Error retrieving provider credentials: {e}") - return None - - async def configure_agent_model( - self, agent_id: str, model_config: AgentModelConfig - ) -> bool: - """Configure model settings for an agent""" - try: - # Validate primary model exists - if model_config.primary_model not in self.available_models: - raise ValueError(f"Unknown model: {model_config.primary_model}") - - # Validate fallback models - for model_id in model_config.fallback_models: - if model_id not in self.available_models: - raise ValueError(f"Unknown fallback model: {model_id}") - - # Store configuration - await DatabaseManager.store_agent_model_config( - agent_id, model_config.__dict__ - ) - - logger.info(f"Configured model settings for agent {agent_id}") - return True - - except Exception as e: - logger.error(f"Error configuring agent model: {e}") - return False - - async def get_agent_model_config(self, agent_id: str) -> Optional[AgentModelConfig]: - """Get model configuration for an agent""" - try: - config_data = await DatabaseManager.get_agent_model_config(agent_id) - - if not config_data: - # Return default configuration - return AgentModelConfig( - agent_id=agent_id, - primary_model="claude-3-5-sonnet-20241022", # Default to Claude 3.5 Sonnet - ) - - return AgentModelConfig(**config_data) - - except Exception as e: - logger.error(f"Error getting agent model config: {e}") - return None - - async def get_model_for_task( - self, - agent_id: str, - task_capabilities: List[ModelCapability], - cost_limit: Optional[float] = None, - ) -> Optional[str]: - """Select best model for a task based on capabilities and cost""" - try: - agent_config = await self.get_agent_model_config(agent_id) - if not agent_config: - return None - - # Check if primary model supports required capabilities - primary_model = self.available_models.get(agent_config.primary_model) - if primary_model and all( - cap in primary_model.capabilities for cap in task_capabilities - ): - # Check cost limit if specified - if ( - cost_limit is None - or self._estimate_task_cost(primary_model, 1000) <= cost_limit - ): - return agent_config.primary_model - - # Try fallback models - for model_id in agent_config.fallback_models: - model = self.available_models.get(model_id) - if model and all( - cap in model.capabilities for cap in task_capabilities - ): - if ( - cost_limit is None - or self._estimate_task_cost(model, 1000) <= cost_limit - ): - return model_id - - # No suitable model found - logger.warning( - f"No suitable model found for agent {agent_id} with capabilities {task_capabilities}" - ) - return None - - except Exception as e: - logger.error(f"Error selecting model for task: {e}") - return None - - def _estimate_task_cost(self, model: ModelSpec, estimated_tokens: int) -> float: - """Estimate cost for a task with given token count""" - # Simple estimation assuming 70% input, 30% output tokens - input_tokens = int(estimated_tokens * 0.7) - output_tokens = int(estimated_tokens * 0.3) - - input_cost = (input_tokens / 1000) * model.cost_per_input_token - output_cost = (output_tokens / 1000) * model.cost_per_output_token - - return input_cost + output_cost - - async def get_available_models( - self, - organization_id: str, - provider: Optional[ModelProvider] = None, - capabilities: Optional[List[ModelCapability]] = None, - ) -> List[Dict[str, Any]]: - """Get available models with provider credential validation""" - available = [] - - for model_id, model in self.available_models.items(): - # Filter by provider if specified - if provider and model.provider != provider: - continue - - # Filter by capabilities if specified - if capabilities and not all( - cap in model.capabilities for cap in capabilities - ): - continue - - # Check if organization has credentials for this provider - has_credentials = ( - await self.get_provider_credentials(organization_id, model.provider) - is not None - ) - - model_info = { - "model_id": model_id, - "provider": model.provider.value, - "name": model.name, - "description": model.description, - "capabilities": [cap.value for cap in model.capabilities], - "context_window": model.context_window, - "max_output_tokens": model.max_output_tokens, - "cost_per_input_token": model.cost_per_input_token, - "cost_per_output_token": model.cost_per_output_token, - "supports_streaming": model.supports_streaming, - "supports_function_calling": model.supports_function_calling, - "supports_vision": model.supports_vision, - "supports_json_mode": model.supports_json_mode, - "recommended_use_cases": model.recommended_use_cases, - "has_credentials": has_credentials, - "available": has_credentials, - } - - available.append(model_info) - - return available - - async def get_organization_model_usage( - self, organization_id: str, days: int = 30 - ) -> Dict[str, Any]: - """Get model usage statistics for an organization""" - try: - usage_data = await DatabaseManager.get_model_usage_stats( - organization_id, days - ) - - return { - "organization_id": organization_id, - "period_days": days, - "total_requests": usage_data.get("total_requests", 0), - "total_tokens": usage_data.get("total_tokens", 0), - "total_cost": usage_data.get("total_cost", 0.0), - "model_breakdown": usage_data.get("model_breakdown", {}), - "agent_breakdown": usage_data.get("agent_breakdown", {}), - "daily_usage": usage_data.get("daily_usage", []), - } - - except Exception as e: - logger.error(f"Error getting model usage: {e}") - return {} - - async def estimate_task_cost( - self, agent_id: str, task_description: str, estimated_complexity: str = "medium" - ) -> Dict[str, Any]: - """Estimate cost for a task execution""" - try: - agent_config = await self.get_agent_model_config(agent_id) - if not agent_config: - return {"error": "Agent configuration not found"} - - model = self.available_models.get(agent_config.primary_model) - if not model: - return {"error": "Model specification not found"} - - # Estimate token usage based on complexity - token_estimates = { - "low": 2000, - "medium": 5000, - "high": 10000, - "very_high": 20000, - } - - estimated_tokens = token_estimates.get(estimated_complexity, 5000) - estimated_cost = self._estimate_task_cost(model, estimated_tokens) - - return { - "agent_id": agent_id, - "model": model.model_id, - "estimated_tokens": estimated_tokens, - "estimated_cost_usd": round(estimated_cost, 4), - "complexity": estimated_complexity, - "cost_breakdown": { - "input_cost": (estimated_tokens * 0.7 / 1000) - * model.cost_per_input_token, - "output_cost": (estimated_tokens * 0.3 / 1000) - * model.cost_per_output_token, - }, - } - - except Exception as e: - logger.error(f"Error estimating task cost: {e}") - return {"error": str(e)} - - -# Global instance -model_config_manager = ModelConfigurationManager() +""" +Model Configuration and API Key Management for FuzeAgent + +Manages model configurations and API keys for different AI providers at the +organization level, with secure storage and agent-specific model selection. + +Supports: +- Multiple AI model providers (Anthropic, OpenAI, Google, etc.) +- Organization-level API key management +- Agent-specific model configurations +- Secure credential storage and access +- Model capability matching +- Cost optimization +""" + +import base64 +import json +import logging +import os +import tempfile +from dataclasses import dataclass, field +from datetime import datetime +from enum import Enum +from typing import Any, Dict, List, Optional, Union + +from cryptography.fernet import Fernet + +from .database import DatabaseManager + +logger = logging.getLogger(__name__) + + +class ModelProvider(str, Enum): + ANTHROPIC = "anthropic" + OPENAI = "openai" + GOOGLE = "google" + AZURE_OPENAI = "azure_openai" + COHERE = "cohere" + HUGGINGFACE = "huggingface" + OLLAMA = "ollama" + CUSTOM = "custom" + + +class ModelCapability(str, Enum): + TEXT_GENERATION = "text_generation" + CODE_GENERATION = "code_generation" + REASONING = "reasoning" + ANALYSIS = "analysis" + CONVERSATION = "conversation" + FUNCTION_CALLING = "function_calling" + VISION = "vision" + EMBEDDINGS = "embeddings" + + +@dataclass +class ModelSpec: + """Specification for an AI model""" + + model_id: str + provider: ModelProvider + name: str + description: str + capabilities: List[ModelCapability] + context_window: int + max_output_tokens: int + cost_per_input_token: float # USD per 1K tokens + cost_per_output_token: float # USD per 1K tokens + supports_streaming: bool = True + supports_function_calling: bool = False + supports_vision: bool = False + supports_json_mode: bool = False + temperature_range: tuple = (0.0, 2.0) + recommended_use_cases: List[str] = field(default_factory=list) + + +@dataclass +class ProviderCredentials: + """Secure storage for provider API credentials""" + + provider: ModelProvider + encrypted_api_key: str + endpoint_url: Optional[str] = None + additional_config: Dict[str, Any] = field(default_factory=dict) + created_at: datetime = field(default_factory=datetime.now) + last_used: Optional[datetime] = None + is_active: bool = True + + +@dataclass +class AgentModelConfig: + """Model configuration for an agent""" + + agent_id: str + primary_model: str # model_id + fallback_models: List[str] = field(default_factory=list) + temperature: float = 0.7 + max_tokens: int = 4096 + top_p: float = 1.0 + frequency_penalty: float = 0.0 + presence_penalty: float = 0.0 + custom_instructions: str = "" + use_function_calling: bool = True + streaming_enabled: bool = True + cost_limit_per_task: Optional[float] = None # USD + created_at: datetime = field(default_factory=datetime.now) + updated_at: datetime = field(default_factory=datetime.now) + + +class ModelConfigurationManager: + """ + Manages AI model configurations and provider credentials. + + Features: + - Secure API key storage with encryption + - Organization-level credential management + - Agent-specific model configurations + - Model capability matching + - Cost tracking and limits + - Automatic fallback handling + """ + + def __init__(self): + self.encryption_key = self._get_or_create_encryption_key() + self.fernet = Fernet(self.encryption_key) + self.available_models = self._initialize_model_specs() + + def _get_or_create_encryption_key(self) -> bytes: + """Get or create encryption key for API credentials""" + # Avoid a hard-coded world-readable /tmp path; allow override and fall + # back to the platform temp dir (still /tmp inside the Linux container). + key_dir = os.getenv("FUZEAGENT_KEY_DIR", tempfile.gettempdir()) + key_file = os.path.join(key_dir, "fuzeagent_encryption.key") + + if os.path.exists(key_file): + with open(key_file, "rb") as f: + return f.read() + else: + key = Fernet.generate_key() + with open(key_file, "wb") as f: + f.write(key) + return key + + def _initialize_model_specs(self) -> Dict[str, ModelSpec]: + """Initialize available model specifications""" + models = {} + + # Anthropic Claude models + models["claude-3-5-sonnet-20241022"] = ModelSpec( + model_id="claude-3-5-sonnet-20241022", + provider=ModelProvider.ANTHROPIC, + name="Claude 3.5 Sonnet", + description="Most intelligent model for complex reasoning and coding", + capabilities=[ + ModelCapability.TEXT_GENERATION, + ModelCapability.CODE_GENERATION, + ModelCapability.REASONING, + ModelCapability.ANALYSIS, + ModelCapability.CONVERSATION, + ModelCapability.FUNCTION_CALLING, + ModelCapability.VISION, + ], + context_window=200000, + max_output_tokens=8192, + cost_per_input_token=0.003, + cost_per_output_token=0.015, + supports_function_calling=True, + supports_vision=True, + supports_json_mode=True, + recommended_use_cases=[ + "complex coding", + "reasoning", + "analysis", + "research", + ], + ) + + models["claude-3-haiku-20240307"] = ModelSpec( + model_id="claude-3-haiku-20240307", + provider=ModelProvider.ANTHROPIC, + name="Claude 3 Haiku", + description="Fastest and most cost-effective model for simple tasks", + capabilities=[ + ModelCapability.TEXT_GENERATION, + ModelCapability.CODE_GENERATION, + ModelCapability.CONVERSATION, + ], + context_window=200000, + max_output_tokens=4096, + cost_per_input_token=0.00025, + cost_per_output_token=0.00125, + supports_function_calling=True, + supports_vision=True, + recommended_use_cases=[ + "simple tasks", + "quick responses", + "cost optimization", + ], + ) + + # OpenAI GPT models + models["gpt-4o"] = ModelSpec( + model_id="gpt-4o", + provider=ModelProvider.OPENAI, + name="GPT-4 Omni", + description="OpenAI's most capable multimodal model", + capabilities=[ + ModelCapability.TEXT_GENERATION, + ModelCapability.CODE_GENERATION, + ModelCapability.REASONING, + ModelCapability.ANALYSIS, + ModelCapability.CONVERSATION, + ModelCapability.FUNCTION_CALLING, + ModelCapability.VISION, + ], + context_window=128000, + max_output_tokens=4096, + cost_per_input_token=0.005, + cost_per_output_token=0.015, + supports_function_calling=True, + supports_vision=True, + supports_json_mode=True, + recommended_use_cases=[ + "multimodal tasks", + "function calling", + "complex reasoning", + ], + ) + + models["gpt-4o-mini"] = ModelSpec( + model_id="gpt-4o-mini", + provider=ModelProvider.OPENAI, + name="GPT-4 Omni Mini", + description="Cost-effective model for simpler tasks", + capabilities=[ + ModelCapability.TEXT_GENERATION, + ModelCapability.CODE_GENERATION, + ModelCapability.CONVERSATION, + ModelCapability.FUNCTION_CALLING, + ], + context_window=128000, + max_output_tokens=16384, + cost_per_input_token=0.00015, + cost_per_output_token=0.0006, + supports_function_calling=True, + supports_json_mode=True, + recommended_use_cases=["cost optimization", "simple tasks", "high volume"], + ) + + # Google Gemini models + models["gemini-1.5-pro"] = ModelSpec( + model_id="gemini-1.5-pro", + provider=ModelProvider.GOOGLE, + name="Gemini 1.5 Pro", + description="Google's most capable model with long context", + capabilities=[ + ModelCapability.TEXT_GENERATION, + ModelCapability.CODE_GENERATION, + ModelCapability.REASONING, + ModelCapability.ANALYSIS, + ModelCapability.CONVERSATION, + ModelCapability.FUNCTION_CALLING, + ModelCapability.VISION, + ], + context_window=2000000, # 2M tokens + max_output_tokens=8192, + cost_per_input_token=0.00125, + cost_per_output_token=0.005, + supports_function_calling=True, + supports_vision=True, + recommended_use_cases=[ + "long context", + "document analysis", + "multimodal tasks", + ], + ) + + return models + + async def store_provider_credentials( + self, + organization_id: str, + provider: ModelProvider, + api_key: str, + endpoint_url: Optional[str] = None, + additional_config: Optional[Dict[str, Any]] = None, + ) -> bool: + """Store encrypted API credentials for a provider""" + try: + # Encrypt the API key + encrypted_key = self.fernet.encrypt(api_key.encode()).decode() + + credentials = ProviderCredentials( + provider=provider, + encrypted_api_key=encrypted_key, + endpoint_url=endpoint_url, + additional_config=additional_config or {}, + ) + + # Store in database + await DatabaseManager.store_provider_credentials( + organization_id, credentials.__dict__ + ) + + logger.info( + f"Stored credentials for {provider} in organization {organization_id}" + ) + return True + + except Exception as e: + logger.error(f"Error storing provider credentials: {e}") + return False + + async def get_provider_credentials( + self, organization_id: str, provider: ModelProvider + ) -> Optional[str]: + """Get decrypted API key for a provider""" + try: + credentials_data = await DatabaseManager.get_provider_credentials( + organization_id, provider.value + ) + + if not credentials_data: + return None + + # Decrypt the API key + encrypted_key = credentials_data["encrypted_api_key"] + decrypted_key = self.fernet.decrypt(encrypted_key.encode()).decode() + + # Update last used timestamp + await DatabaseManager.update_credentials_last_used( + organization_id, provider.value + ) + + return decrypted_key + + except Exception as e: + logger.error(f"Error retrieving provider credentials: {e}") + return None + + async def configure_agent_model( + self, agent_id: str, model_config: AgentModelConfig + ) -> bool: + """Configure model settings for an agent""" + try: + # Validate primary model exists + if model_config.primary_model not in self.available_models: + raise ValueError(f"Unknown model: {model_config.primary_model}") + + # Validate fallback models + for model_id in model_config.fallback_models: + if model_id not in self.available_models: + raise ValueError(f"Unknown fallback model: {model_id}") + + # Store configuration + await DatabaseManager.store_agent_model_config( + agent_id, model_config.__dict__ + ) + + logger.info(f"Configured model settings for agent {agent_id}") + return True + + except Exception as e: + logger.error(f"Error configuring agent model: {e}") + return False + + async def get_agent_model_config(self, agent_id: str) -> Optional[AgentModelConfig]: + """Get model configuration for an agent""" + try: + config_data = await DatabaseManager.get_agent_model_config(agent_id) + + if not config_data: + # Return default configuration + return AgentModelConfig( + agent_id=agent_id, + primary_model="claude-3-5-sonnet-20241022", # Default to Claude 3.5 Sonnet + ) + + return AgentModelConfig(**config_data) + + except Exception as e: + logger.error(f"Error getting agent model config: {e}") + return None + + async def get_model_for_task( + self, + agent_id: str, + task_capabilities: List[ModelCapability], + cost_limit: Optional[float] = None, + ) -> Optional[str]: + """Select best model for a task based on capabilities and cost""" + try: + agent_config = await self.get_agent_model_config(agent_id) + if not agent_config: + return None + + # Check if primary model supports required capabilities + primary_model = self.available_models.get(agent_config.primary_model) + if primary_model and all( + cap in primary_model.capabilities for cap in task_capabilities + ): + # Check cost limit if specified + if ( + cost_limit is None + or self._estimate_task_cost(primary_model, 1000) <= cost_limit + ): + return agent_config.primary_model + + # Try fallback models + for model_id in agent_config.fallback_models: + model = self.available_models.get(model_id) + if model and all( + cap in model.capabilities for cap in task_capabilities + ): + if ( + cost_limit is None + or self._estimate_task_cost(model, 1000) <= cost_limit + ): + return model_id + + # No suitable model found + logger.warning( + f"No suitable model found for agent {agent_id} with capabilities {task_capabilities}" + ) + return None + + except Exception as e: + logger.error(f"Error selecting model for task: {e}") + return None + + def _estimate_task_cost(self, model: ModelSpec, estimated_tokens: int) -> float: + """Estimate cost for a task with given token count""" + # Simple estimation assuming 70% input, 30% output tokens + input_tokens = int(estimated_tokens * 0.7) + output_tokens = int(estimated_tokens * 0.3) + + input_cost = (input_tokens / 1000) * model.cost_per_input_token + output_cost = (output_tokens / 1000) * model.cost_per_output_token + + return input_cost + output_cost + + async def get_available_models( + self, + organization_id: str, + provider: Optional[ModelProvider] = None, + capabilities: Optional[List[ModelCapability]] = None, + ) -> List[Dict[str, Any]]: + """Get available models with provider credential validation""" + available = [] + + for model_id, model in self.available_models.items(): + # Filter by provider if specified + if provider and model.provider != provider: + continue + + # Filter by capabilities if specified + if capabilities and not all( + cap in model.capabilities for cap in capabilities + ): + continue + + # Check if organization has credentials for this provider + has_credentials = ( + await self.get_provider_credentials(organization_id, model.provider) + is not None + ) + + model_info = { + "model_id": model_id, + "provider": model.provider.value, + "name": model.name, + "description": model.description, + "capabilities": [cap.value for cap in model.capabilities], + "context_window": model.context_window, + "max_output_tokens": model.max_output_tokens, + "cost_per_input_token": model.cost_per_input_token, + "cost_per_output_token": model.cost_per_output_token, + "supports_streaming": model.supports_streaming, + "supports_function_calling": model.supports_function_calling, + "supports_vision": model.supports_vision, + "supports_json_mode": model.supports_json_mode, + "recommended_use_cases": model.recommended_use_cases, + "has_credentials": has_credentials, + "available": has_credentials, + } + + available.append(model_info) + + return available + + async def get_organization_model_usage( + self, organization_id: str, days: int = 30 + ) -> Dict[str, Any]: + """Get model usage statistics for an organization""" + try: + usage_data = await DatabaseManager.get_model_usage_stats( + organization_id, days + ) + + return { + "organization_id": organization_id, + "period_days": days, + "total_requests": usage_data.get("total_requests", 0), + "total_tokens": usage_data.get("total_tokens", 0), + "total_cost": usage_data.get("total_cost", 0.0), + "model_breakdown": usage_data.get("model_breakdown", {}), + "agent_breakdown": usage_data.get("agent_breakdown", {}), + "daily_usage": usage_data.get("daily_usage", []), + } + + except Exception as e: + logger.error(f"Error getting model usage: {e}") + return {} + + async def estimate_task_cost( + self, agent_id: str, task_description: str, estimated_complexity: str = "medium" + ) -> Dict[str, Any]: + """Estimate cost for a task execution""" + try: + agent_config = await self.get_agent_model_config(agent_id) + if not agent_config: + return {"error": "Agent configuration not found"} + + model = self.available_models.get(agent_config.primary_model) + if not model: + return {"error": "Model specification not found"} + + # Estimate token usage based on complexity + token_estimates = { + "low": 2000, + "medium": 5000, + "high": 10000, + "very_high": 20000, + } + + estimated_tokens = token_estimates.get(estimated_complexity, 5000) + estimated_cost = self._estimate_task_cost(model, estimated_tokens) + + return { + "agent_id": agent_id, + "model": model.model_id, + "estimated_tokens": estimated_tokens, + "estimated_cost_usd": round(estimated_cost, 4), + "complexity": estimated_complexity, + "cost_breakdown": { + "input_cost": (estimated_tokens * 0.7 / 1000) + * model.cost_per_input_token, + "output_cost": (estimated_tokens * 0.3 / 1000) + * model.cost_per_output_token, + }, + } + + except Exception as e: + logger.error(f"Error estimating task cost: {e}") + return {"error": str(e)} + + +# Global instance +model_config_manager = ModelConfigurationManager() diff --git a/services/orchestrator/multi_agent_coordinator.py b/services/orchestrator/multi_agent_coordinator.py index 3c80db9..f1951ee 100644 --- a/services/orchestrator/multi_agent_coordinator.py +++ b/services/orchestrator/multi_agent_coordinator.py @@ -1,938 +1,938 @@ -""" -Multi-Agent Coordination System for FuzeAgent - -Enables multiple agents to collaborate on complex tasks through: -- Task decomposition and delegation -- Agent communication and synchronization -- Dependency management -- Result aggregation -- Conflict resolution -- Progress monitoring across agent teams - -This system allows for autonomous coordination of development teams -where agents can request help, delegate subtasks, and coordinate -work without human intervention. -""" - -import asyncio -import json -import logging -import uuid -from dataclasses import dataclass, field -from datetime import datetime, timedelta -from enum import Enum -from typing import Any, Dict, List, Optional, Set, Tuple - -from .database import DatabaseManager -from .task_execution_engine import TaskExecutionEngine, TaskStatus - -logger = logging.getLogger(__name__) - - -class CoordinationMode(str, Enum): - SEQUENTIAL = "sequential" # Tasks executed one after another - PARALLEL = "parallel" # Tasks executed simultaneously - HIERARCHICAL = "hierarchical" # Manager delegates to subordinates - COLLABORATIVE = "collaborative" # Agents work together on shared task - - -class AgentRole(str, Enum): - COORDINATOR = "coordinator" # Leads the coordination - PARTICIPANT = "participant" # Participates in coordination - OBSERVER = "observer" # Observes but doesn't execute - - -class CoordinationStatus(str, Enum): - INITIALIZING = "initializing" - PLANNING = "planning" - EXECUTING = "executing" - SYNCHRONIZING = "synchronizing" - REVIEWING = "reviewing" - COMPLETED = "completed" - FAILED = "failed" - CANCELLED = "cancelled" - - -@dataclass -class AgentCapability: - """Represents an agent's capability""" - - skill: str - proficiency: float # 0.0 to 1.0 - availability: bool - current_load: float # 0.0 to 1.0 - - -@dataclass -class TaskDependency: - """Represents a dependency between tasks""" - - dependent_task_id: str - prerequisite_task_id: str - dependency_type: str # "blocking", "soft", "informational" - - -@dataclass -class CoordinationPlan: - """Represents a plan for multi-agent coordination""" - - plan_id: str - root_task_id: str - coordination_mode: CoordinationMode - participating_agents: List[str] - task_assignments: Dict[str, str] # task_id -> agent_id - dependencies: List[TaskDependency] - estimated_completion: datetime - created_at: datetime - - -@dataclass -class AgentCommunication: - """Represents communication between agents""" - - communication_id: str - from_agent_id: str - to_agent_id: str - message_type: str # "request", "response", "notification", "question" - content: str - metadata: Dict[str, Any] - timestamp: datetime - response_id: Optional[str] = None - - -@dataclass -class CoordinationSession: - """Represents an active multi-agent coordination session""" - - session_id: str - root_task_id: str - coordinator_agent_id: str - participating_agents: Set[str] - coordination_mode: CoordinationMode - status: CoordinationStatus - plan: Optional[CoordinationPlan] - communications: List[AgentCommunication] = field(default_factory=list) - subtasks: Dict[str, str] = field(default_factory=dict) # subtask_id -> agent_id - started_at: datetime = field(default_factory=datetime.now) - completed_at: Optional[datetime] = None - result: Optional[Dict[str, Any]] = None - - -class MultiAgentCoordinator: - """ - Orchestrates multi-agent collaboration for complex tasks. - - Features: - - Automatic task decomposition - - Agent capability matching - - Dynamic load balancing - - Inter-agent communication - - Dependency resolution - - Progress synchronization - - Conflict resolution - """ - - def __init__(self, task_execution_engine: TaskExecutionEngine): - self.task_execution_engine = task_execution_engine - self.active_sessions: Dict[str, CoordinationSession] = {} - self.agent_capabilities: Dict[str, List[AgentCapability]] = {} - self.running = False - self.coordination_workers: List[asyncio.Task] = [] - - # Configuration - self.max_concurrent_coordinations = 10 - self.communication_timeout = 300 # 5 minutes - self.synchronization_interval = 30 # 30 seconds - - async def start(self): - """Start the multi-agent coordinator""" - logger.info("Starting MultiAgentCoordinator") - self.running = True - - # Start coordination workers - self.coordination_workers = [ - asyncio.create_task(self._coordination_worker()), - asyncio.create_task(self._communication_worker()), - asyncio.create_task(self._synchronization_worker()), - ] - - logger.info("MultiAgentCoordinator started") - - async def stop(self): - """Stop the multi-agent coordinator""" - logger.info("Stopping MultiAgentCoordinator") - self.running = False - - # Cancel workers - for worker in self.coordination_workers: - worker.cancel() - - try: - await asyncio.gather(*self.coordination_workers, return_exceptions=True) - except Exception as e: - logger.error(f"Error stopping coordination workers: {e}") - - # Clean up active sessions - for session_id in list(self.active_sessions.keys()): - await self._cleanup_session(session_id) - - logger.info("MultiAgentCoordinator stopped") - - async def initiate_coordination( - self, - task_id: str, - coordination_mode: CoordinationMode = CoordinationMode.COLLABORATIVE, - required_agents: Optional[List[str]] = None, - required_skills: Optional[List[str]] = None, - ) -> str: - """ - Initiate multi-agent coordination for a complex task. - - Args: - task_id: The root task to coordinate - coordination_mode: How agents should coordinate - required_agents: Specific agents to include - required_skills: Required skills for the task - - Returns: - Coordination session ID - """ - logger.info(f"Initiating coordination for task {task_id}") - - try: - # Get task information - task_data = await DatabaseManager.get_task(task_id) - if not task_data: - raise ValueError(f"Task {task_id} not found") - - # Analyze task complexity and determine if coordination is needed - complexity_analysis = await self._analyze_task_complexity(task_data) - - if not complexity_analysis["requires_coordination"]: - logger.info(f"Task {task_id} does not require coordination") - return None - - # Find suitable agents - if required_agents: - selected_agents = required_agents - else: - selected_agents = await self._select_agents_for_task( - task_data, required_skills, complexity_analysis - ) - - if len(selected_agents) < 2: - logger.warning( - f"Not enough agents available for coordination: {len(selected_agents)}" - ) - return None - - # Determine coordinator agent (first agent or most experienced) - coordinator_agent_id = await self._select_coordinator( - selected_agents, task_data - ) - - # Create coordination session - session_id = str(uuid.uuid4()) - session = CoordinationSession( - session_id=session_id, - root_task_id=task_id, - coordinator_agent_id=coordinator_agent_id, - participating_agents=set(selected_agents), - coordination_mode=coordination_mode, - status=CoordinationStatus.INITIALIZING, - ) - - self.active_sessions[session_id] = session - - logger.info( - f"Created coordination session {session_id} with {len(selected_agents)} agents" - ) - return session_id - - except Exception as e: - logger.error(f"Failed to initiate coordination for task {task_id}: {e}") - raise - - async def get_coordination_status( - self, session_id: str - ) -> Optional[Dict[str, Any]]: - """Get status of a coordination session""" - session = self.active_sessions.get(session_id) - if not session: - return None - - # Get subtask statuses - subtask_statuses = {} - for subtask_id, agent_id in session.subtasks.items(): - status = await self.task_execution_engine.get_execution_status(subtask_id) - subtask_statuses[subtask_id] = { - "agent_id": agent_id, - "status": status.get("status", "unknown") if status else "unknown", - } - - return { - "session_id": session_id, - "root_task_id": session.root_task_id, - "coordinator": session.coordinator_agent_id, - "participating_agents": list(session.participating_agents), - "coordination_mode": session.coordination_mode.value, - "status": session.status.value, - "subtasks": subtask_statuses, - "communications_count": len(session.communications), - "started_at": session.started_at.isoformat(), - "completed_at": ( - session.completed_at.isoformat() if session.completed_at else None - ), - "estimated_completion": ( - session.plan.estimated_completion.isoformat() if session.plan else None - ), - } - - async def send_agent_communication( - self, - from_agent_id: str, - to_agent_id: str, - message_type: str, - content: str, - metadata: Optional[Dict[str, Any]] = None, - ) -> str: - """Send communication between agents""" - - communication_id = str(uuid.uuid4()) - communication = AgentCommunication( - communication_id=communication_id, - from_agent_id=from_agent_id, - to_agent_id=to_agent_id, - message_type=message_type, - content=content, - metadata=metadata or {}, - timestamp=datetime.now(), - ) - - # Find coordination session for these agents - session = self._find_session_by_agents([from_agent_id, to_agent_id]) - if session: - session.communications.append(communication) - - # Store in database - await self._store_communication(communication) - - logger.info( - f"Agent communication {communication_id}: {from_agent_id} -> {to_agent_id}" - ) - return communication_id - - async def cancel_coordination(self, session_id: str) -> bool: - """Cancel an active coordination session""" - session = self.active_sessions.get(session_id) - if not session: - return False - - try: - # Cancel all subtasks - for subtask_id in session.subtasks.keys(): - await self.task_execution_engine.cancel_task_execution(subtask_id) - - # Update session status - session.status = CoordinationStatus.CANCELLED - session.completed_at = datetime.now() - - # Cleanup - await self._cleanup_session(session_id) - - logger.info(f"Cancelled coordination session {session_id}") - return True - - except Exception as e: - logger.error(f"Error cancelling coordination session {session_id}: {e}") - return False - - # Private methods - - async def _coordination_worker(self): - """Main coordination worker that manages session lifecycle""" - while self.running: - try: - # Process sessions that need attention - for session_id, session in list(self.active_sessions.items()): - try: - await self._process_coordination_session(session) - except Exception as e: - logger.error( - f"Error processing coordination session {session_id}: {e}" - ) - - await asyncio.sleep(5) # Check every 5 seconds - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Error in coordination worker: {e}") - await asyncio.sleep(10) - - async def _communication_worker(self): - """Worker that handles inter-agent communications""" - while self.running: - try: - # Process pending communications - for session in self.active_sessions.values(): - for comm in session.communications: - if not comm.response_id and comm.message_type == "request": - # Check if communication has timed out - if ( - datetime.now() - comm.timestamp - ).total_seconds() > self.communication_timeout: - await self._handle_communication_timeout(session, comm) - - await asyncio.sleep(10) # Check every 10 seconds - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Error in communication worker: {e}") - await asyncio.sleep(10) - - async def _synchronization_worker(self): - """Worker that synchronizes coordination sessions""" - while self.running: - try: - for session_id, session in list(self.active_sessions.items()): - if session.status == CoordinationStatus.EXECUTING: - await self._synchronize_session(session) - - await asyncio.sleep(self.synchronization_interval) - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Error in synchronization worker: {e}") - await asyncio.sleep(self.synchronization_interval) - - async def _process_coordination_session(self, session: CoordinationSession): - """Process a coordination session based on its current status""" - - if session.status == CoordinationStatus.INITIALIZING: - await self._initialize_session(session) - elif session.status == CoordinationStatus.PLANNING: - await self._plan_coordination(session) - elif session.status == CoordinationStatus.EXECUTING: - await self._monitor_execution(session) - elif session.status == CoordinationStatus.REVIEWING: - await self._review_coordination(session) - - async def _initialize_session(self, session: CoordinationSession): - """Initialize a coordination session""" - try: - # Get detailed task information - task_data = await DatabaseManager.get_task(session.root_task_id) - - # Update agent capabilities - await self._update_agent_capabilities(list(session.participating_agents)) - - # Move to planning phase - session.status = CoordinationStatus.PLANNING - - logger.info(f"Initialized coordination session {session.session_id}") - - except Exception as e: - logger.error(f"Error initializing session {session.session_id}: {e}") - session.status = CoordinationStatus.FAILED - - async def _plan_coordination(self, session: CoordinationSession): - """Create coordination plan""" - try: - # Get task data - task_data = await DatabaseManager.get_task(session.root_task_id) - - # Decompose task into subtasks - subtasks = await self._decompose_task(task_data, session.coordination_mode) - - # Assign agents to subtasks - assignments = await self._assign_agents_to_subtasks( - subtasks, list(session.participating_agents) - ) - - # Create dependencies - dependencies = await self._create_task_dependencies( - subtasks, session.coordination_mode - ) - - # Estimate completion time - estimated_completion = await self._estimate_coordination_completion( - subtasks, assignments, dependencies - ) - - # Create coordination plan - plan = CoordinationPlan( - plan_id=str(uuid.uuid4()), - root_task_id=session.root_task_id, - coordination_mode=session.coordination_mode, - participating_agents=list(session.participating_agents), - task_assignments=assignments, - dependencies=dependencies, - estimated_completion=estimated_completion, - created_at=datetime.now(), - ) - - session.plan = plan - session.status = CoordinationStatus.EXECUTING - - # Create subtasks in database and start execution - for subtask_data in subtasks: - subtask_id = await self._create_subtask(subtask_data, assignments) - session.subtasks[subtask_id] = assignments[subtask_data["id"]] - - # Start subtask execution - await self.task_execution_engine.start_task_execution(subtask_id) - - logger.info( - f"Created coordination plan for session {session.session_id} with {len(subtasks)} subtasks" - ) - - except Exception as e: - logger.error(f"Error planning coordination {session.session_id}: {e}") - session.status = CoordinationStatus.FAILED - - async def _monitor_execution(self, session: CoordinationSession): - """Monitor execution of coordinated tasks""" - try: - # Check status of all subtasks - completed_subtasks = 0 - failed_subtasks = 0 - - for subtask_id in session.subtasks.keys(): - status = await self.task_execution_engine.get_execution_status( - subtask_id - ) - if status: - if status.get("status") == "completed": - completed_subtasks += 1 - elif status.get("status") == "failed": - failed_subtasks += 1 - - total_subtasks = len(session.subtasks) - - # Check if coordination is complete - if completed_subtasks == total_subtasks: - session.status = CoordinationStatus.REVIEWING - elif failed_subtasks > 0: - # Handle failures - await self._handle_coordination_failures(session) - - except Exception as e: - logger.error(f"Error monitoring execution {session.session_id}: {e}") - - async def _review_coordination(self, session: CoordinationSession): - """Review completed coordination and aggregate results""" - try: - # Collect results from all subtasks - results = {} - for subtask_id, agent_id in session.subtasks.items(): - status = await self.task_execution_engine.get_execution_status( - subtask_id - ) - if status: - results[subtask_id] = { - "agent_id": agent_id, - "status": status.get("status"), - "result": status.get("result"), - } - - # Aggregate results - coordination_result = await self._aggregate_coordination_results(results) - - # Complete coordination - session.status = CoordinationStatus.COMPLETED - session.completed_at = datetime.now() - session.result = coordination_result - - # Update root task status - await DatabaseManager.update_task_status( - session.root_task_id, "completed", coordination_result - ) - - logger.info(f"Completed coordination session {session.session_id}") - - # Schedule cleanup - asyncio.create_task(self._cleanup_session(session.session_id)) - - except Exception as e: - logger.error(f"Error reviewing coordination {session.session_id}: {e}") - session.status = CoordinationStatus.FAILED - - async def _analyze_task_complexity( - self, task_data: Dict[str, Any] - ) -> Dict[str, Any]: - """Analyze task complexity to determine if coordination is needed""" - - description = task_data.get("description", "") - title = task_data.get("title", "") - - # Simple heuristics for complexity analysis - complexity_indicators = [ - "multiple components" in description.lower(), - "frontend and backend" in description.lower(), - "database and api" in description.lower(), - "testing and deployment" in description.lower(), - len(description.split()) > 50, # Long description - "integrate" in description.lower(), - "coordinate" in description.lower(), - "collaborate" in description.lower(), - ] - - complexity_score = sum(complexity_indicators) / len(complexity_indicators) - - return { - "requires_coordination": complexity_score > 0.3, - "complexity_score": complexity_score, - "estimated_agents_needed": min(max(2, int(complexity_score * 5)), 5), - "estimated_duration_hours": max(4, int(complexity_score * 24)), - } - - async def _select_agents_for_task( - self, - task_data: Dict[str, Any], - required_skills: Optional[List[str]], - complexity_analysis: Dict[str, Any], - ) -> List[str]: - """Select appropriate agents for the task""" - - # Get all available agents - all_agents = await DatabaseManager.get_all_agents() - - # Filter by availability and skills - suitable_agents = [] - for agent in all_agents: - if agent["status"] == "available": - agent_skills = agent.get("config", {}).get("tools", []) - - # Check skill match - if required_skills: - skill_match = any( - skill in agent_skills for skill in required_skills - ) - else: - skill_match = True - - if skill_match: - suitable_agents.append(agent["id"]) - - # Select optimal number of agents - max_agents = complexity_analysis.get("estimated_agents_needed", 3) - return suitable_agents[:max_agents] - - async def _select_coordinator( - self, agents: List[str], task_data: Dict[str, Any] - ) -> str: - """Select the coordinator agent from available agents""" - - # For now, select the first agent as coordinator - # In production, this would consider agent experience, current load, etc. - return agents[0] - - async def _decompose_task( - self, task_data: Dict[str, Any], mode: CoordinationMode - ) -> List[Dict[str, Any]]: - """Decompose a complex task into subtasks""" - - # Simple task decomposition based on common patterns - description = task_data.get("description", "") - title = task_data.get("title", "") - - subtasks = [] - - # Common subtask patterns - if "frontend" in description.lower() or "ui" in description.lower(): - subtasks.append( - { - "id": f"frontend-{uuid.uuid4()}", - "title": f"Frontend Implementation - {title}", - "description": f"Implement frontend components for: {description}", - "type": "frontend_development", - "estimated_hours": 4, - } - ) - - if "backend" in description.lower() or "api" in description.lower(): - subtasks.append( - { - "id": f"backend-{uuid.uuid4()}", - "title": f"Backend Implementation - {title}", - "description": f"Implement backend services for: {description}", - "type": "backend_development", - "estimated_hours": 6, - } - ) - - if "database" in description.lower() or "data" in description.lower(): - subtasks.append( - { - "id": f"database-{uuid.uuid4()}", - "title": f"Database Design - {title}", - "description": f"Design and implement database schema for: {description}", - "type": "database_development", - "estimated_hours": 3, - } - ) - - if "test" in description.lower(): - subtasks.append( - { - "id": f"testing-{uuid.uuid4()}", - "title": f"Testing - {title}", - "description": f"Create comprehensive tests for: {description}", - "type": "testing", - "estimated_hours": 4, - } - ) - - # If no specific subtasks identified, create generic subtasks - if not subtasks: - subtasks = [ - { - "id": f"analysis-{uuid.uuid4()}", - "title": f"Analysis - {title}", - "description": f"Analyze requirements for: {description}", - "type": "analysis", - "estimated_hours": 2, - }, - { - "id": f"implementation-{uuid.uuid4()}", - "title": f"Implementation - {title}", - "description": f"Implement solution for: {description}", - "type": "implementation", - "estimated_hours": 6, - }, - { - "id": f"review-{uuid.uuid4()}", - "title": f"Review - {title}", - "description": f"Review and validate solution for: {description}", - "type": "review", - "estimated_hours": 2, - }, - ] - - return subtasks - - async def _assign_agents_to_subtasks( - self, subtasks: List[Dict[str, Any]], agents: List[str] - ) -> Dict[str, str]: - """Assign agents to subtasks based on capabilities""" - - assignments = {} - - # Get agent capabilities - agent_data = {} - for agent_id in agents: - agent = await DatabaseManager.get_agent(agent_id) - if agent: - agent_data[agent_id] = agent - - # Simple assignment based on agent type - for subtask in subtasks: - subtask_type = subtask.get("type", "") - best_agent = None - - # Match agent type to subtask type - for agent_id, agent in agent_data.items(): - agent_type = agent.get("type", "") - - if subtask_type.startswith("frontend") and "frontend" in agent_type: - best_agent = agent_id - break - elif subtask_type.startswith("backend") and "backend" in agent_type: - best_agent = agent_id - break - elif subtask_type.startswith("database") and "backend" in agent_type: - best_agent = agent_id - break - elif subtask_type == "testing" and "qa" in agent_type: - best_agent = agent_id - break - - # Fallback to first available agent - if not best_agent: - best_agent = agents[0] - - assignments[subtask["id"]] = best_agent - - return assignments - - async def _create_task_dependencies( - self, subtasks: List[Dict[str, Any]], mode: CoordinationMode - ) -> List[TaskDependency]: - """Create dependencies between subtasks""" - - dependencies = [] - - if mode == CoordinationMode.SEQUENTIAL: - # Create sequential dependencies - for i in range(1, len(subtasks)): - dependencies.append( - TaskDependency( - dependent_task_id=subtasks[i]["id"], - prerequisite_task_id=subtasks[i - 1]["id"], - dependency_type="blocking", - ) - ) - - elif mode == CoordinationMode.HIERARCHICAL: - # Analysis task should complete before implementation - analysis_tasks = [t for t in subtasks if "analysis" in t["type"]] - implementation_tasks = [ - t for t in subtasks if "implementation" in t["type"] - ] - - for impl_task in implementation_tasks: - for analysis_task in analysis_tasks: - dependencies.append( - TaskDependency( - dependent_task_id=impl_task["id"], - prerequisite_task_id=analysis_task["id"], - dependency_type="blocking", - ) - ) - - # PARALLEL and COLLABORATIVE modes have no strict dependencies - - return dependencies - - async def _estimate_coordination_completion( - self, - subtasks: List[Dict[str, Any]], - assignments: Dict[str, str], - dependencies: List[TaskDependency], - ) -> datetime: - """Estimate when coordination will complete""" - - if not dependencies: - # Parallel execution - completion time is max of all subtasks - max_hours = max(subtask.get("estimated_hours", 4) for subtask in subtasks) - else: - # Sequential/dependent execution - sum of critical path - total_hours = sum(subtask.get("estimated_hours", 4) for subtask in subtasks) - max_hours = min(total_hours, 24) # Cap at 24 hours - - return datetime.now() + timedelta(hours=max_hours) - - async def _create_subtask( - self, subtask_data: Dict[str, Any], assignments: Dict[str, str] - ) -> str: - """Create a subtask in the database""" - - subtask_id = str(uuid.uuid4()) - agent_id = assignments.get(subtask_data["id"]) - - # Create task in database - await DatabaseManager.create_task( - task_id=subtask_id, - title=subtask_data["title"], - description=subtask_data["description"], - assigned_to=agent_id, - priority="medium", - metadata={ - "coordination_subtask": True, - "parent_task_type": subtask_data.get("type"), - "estimated_hours": subtask_data.get("estimated_hours", 4), - }, - ) - - return subtask_id - - def _find_session_by_agents( - self, agent_ids: List[str] - ) -> Optional[CoordinationSession]: - """Find coordination session that includes the specified agents""" - for session in self.active_sessions.values(): - if any(agent_id in session.participating_agents for agent_id in agent_ids): - return session - return None - - async def _store_communication(self, communication: AgentCommunication): - """Store agent communication in database""" - # This would store the communication in the database - # For now, just log it - logger.info( - f"Agent communication: {communication.from_agent_id} -> {communication.to_agent_id}: {communication.content}" - ) - - async def _handle_communication_timeout( - self, session: CoordinationSession, communication: AgentCommunication - ): - """Handle communication timeout""" - logger.warning( - f"Communication timeout in session {session.session_id}: {communication.communication_id}" - ) - - # Could implement retry logic or escalation here - - async def _synchronize_session(self, session: CoordinationSession): - """Synchronize coordination session state""" - # Check if any agents need help or coordination - # Update session status based on subtask progress - # Handle any conflicts or issues - pass - - async def _handle_coordination_failures(self, session: CoordinationSession): - """Handle failures in coordination""" - logger.warning( - f"Handling failures in coordination session {session.session_id}" - ) - - # Could implement retry logic, reassignment, or escalation - session.status = CoordinationStatus.FAILED - - async def _aggregate_coordination_results( - self, results: Dict[str, Dict[str, Any]] - ) -> Dict[str, Any]: - """Aggregate results from all coordinated subtasks""" - - successful_subtasks = [ - r for r in results.values() if r["status"] == "completed" - ] - - return { - "coordination_completed": True, - "total_subtasks": len(results), - "successful_subtasks": len(successful_subtasks), - "failed_subtasks": len(results) - len(successful_subtasks), - "results": results, - "completion_time": datetime.now().isoformat(), - } - - async def _update_agent_capabilities(self, agent_ids: List[str]): - """Update cached agent capabilities""" - for agent_id in agent_ids: - agent_data = await DatabaseManager.get_agent(agent_id) - if agent_data: - # Extract capabilities from agent configuration - tools = agent_data.get("config", {}).get("tools", []) - capabilities = [ - AgentCapability( - skill=tool, - proficiency=0.8, # Default proficiency - availability=agent_data.get("status") == "available", - current_load=0.5, # Default load - ) - for tool in tools - ] - self.agent_capabilities[agent_id] = capabilities - - async def _cleanup_session(self, session_id: str): - """Clean up completed coordination session""" - session = self.active_sessions.pop(session_id, None) - if session: - logger.info(f"Cleaned up coordination session {session_id}") - - -# Integration with existing TaskExecutionEngine -def integrate_multi_agent_coordination( - task_execution_engine: TaskExecutionEngine, -) -> MultiAgentCoordinator: - """Create and integrate multi-agent coordinator with task execution engine""" - coordinator = MultiAgentCoordinator(task_execution_engine) - - # Add coordination capabilities to task execution engine - task_execution_engine.multi_agent_coordinator = coordinator - - return coordinator +""" +Multi-Agent Coordination System for FuzeAgent + +Enables multiple agents to collaborate on complex tasks through: +- Task decomposition and delegation +- Agent communication and synchronization +- Dependency management +- Result aggregation +- Conflict resolution +- Progress monitoring across agent teams + +This system allows for autonomous coordination of development teams +where agents can request help, delegate subtasks, and coordinate +work without human intervention. +""" + +import asyncio +import json +import logging +import uuid +from dataclasses import dataclass, field +from datetime import datetime, timedelta +from enum import Enum +from typing import Any, Dict, List, Optional, Set, Tuple + +from .database import DatabaseManager +from .task_execution_engine import TaskExecutionEngine, TaskStatus + +logger = logging.getLogger(__name__) + + +class CoordinationMode(str, Enum): + SEQUENTIAL = "sequential" # Tasks executed one after another + PARALLEL = "parallel" # Tasks executed simultaneously + HIERARCHICAL = "hierarchical" # Manager delegates to subordinates + COLLABORATIVE = "collaborative" # Agents work together on shared task + + +class AgentRole(str, Enum): + COORDINATOR = "coordinator" # Leads the coordination + PARTICIPANT = "participant" # Participates in coordination + OBSERVER = "observer" # Observes but doesn't execute + + +class CoordinationStatus(str, Enum): + INITIALIZING = "initializing" + PLANNING = "planning" + EXECUTING = "executing" + SYNCHRONIZING = "synchronizing" + REVIEWING = "reviewing" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +@dataclass +class AgentCapability: + """Represents an agent's capability""" + + skill: str + proficiency: float # 0.0 to 1.0 + availability: bool + current_load: float # 0.0 to 1.0 + + +@dataclass +class TaskDependency: + """Represents a dependency between tasks""" + + dependent_task_id: str + prerequisite_task_id: str + dependency_type: str # "blocking", "soft", "informational" + + +@dataclass +class CoordinationPlan: + """Represents a plan for multi-agent coordination""" + + plan_id: str + root_task_id: str + coordination_mode: CoordinationMode + participating_agents: List[str] + task_assignments: Dict[str, str] # task_id -> agent_id + dependencies: List[TaskDependency] + estimated_completion: datetime + created_at: datetime + + +@dataclass +class AgentCommunication: + """Represents communication between agents""" + + communication_id: str + from_agent_id: str + to_agent_id: str + message_type: str # "request", "response", "notification", "question" + content: str + metadata: Dict[str, Any] + timestamp: datetime + response_id: Optional[str] = None + + +@dataclass +class CoordinationSession: + """Represents an active multi-agent coordination session""" + + session_id: str + root_task_id: str + coordinator_agent_id: str + participating_agents: Set[str] + coordination_mode: CoordinationMode + status: CoordinationStatus + plan: Optional[CoordinationPlan] + communications: List[AgentCommunication] = field(default_factory=list) + subtasks: Dict[str, str] = field(default_factory=dict) # subtask_id -> agent_id + started_at: datetime = field(default_factory=datetime.now) + completed_at: Optional[datetime] = None + result: Optional[Dict[str, Any]] = None + + +class MultiAgentCoordinator: + """ + Orchestrates multi-agent collaboration for complex tasks. + + Features: + - Automatic task decomposition + - Agent capability matching + - Dynamic load balancing + - Inter-agent communication + - Dependency resolution + - Progress synchronization + - Conflict resolution + """ + + def __init__(self, task_execution_engine: TaskExecutionEngine): + self.task_execution_engine = task_execution_engine + self.active_sessions: Dict[str, CoordinationSession] = {} + self.agent_capabilities: Dict[str, List[AgentCapability]] = {} + self.running = False + self.coordination_workers: List[asyncio.Task] = [] + + # Configuration + self.max_concurrent_coordinations = 10 + self.communication_timeout = 300 # 5 minutes + self.synchronization_interval = 30 # 30 seconds + + async def start(self): + """Start the multi-agent coordinator""" + logger.info("Starting MultiAgentCoordinator") + self.running = True + + # Start coordination workers + self.coordination_workers = [ + asyncio.create_task(self._coordination_worker()), + asyncio.create_task(self._communication_worker()), + asyncio.create_task(self._synchronization_worker()), + ] + + logger.info("MultiAgentCoordinator started") + + async def stop(self): + """Stop the multi-agent coordinator""" + logger.info("Stopping MultiAgentCoordinator") + self.running = False + + # Cancel workers + for worker in self.coordination_workers: + worker.cancel() + + try: + await asyncio.gather(*self.coordination_workers, return_exceptions=True) + except Exception as e: + logger.error(f"Error stopping coordination workers: {e}") + + # Clean up active sessions + for session_id in list(self.active_sessions.keys()): + await self._cleanup_session(session_id) + + logger.info("MultiAgentCoordinator stopped") + + async def initiate_coordination( + self, + task_id: str, + coordination_mode: CoordinationMode = CoordinationMode.COLLABORATIVE, + required_agents: Optional[List[str]] = None, + required_skills: Optional[List[str]] = None, + ) -> str: + """ + Initiate multi-agent coordination for a complex task. + + Args: + task_id: The root task to coordinate + coordination_mode: How agents should coordinate + required_agents: Specific agents to include + required_skills: Required skills for the task + + Returns: + Coordination session ID + """ + logger.info(f"Initiating coordination for task {task_id}") + + try: + # Get task information + task_data = await DatabaseManager.get_task(task_id) + if not task_data: + raise ValueError(f"Task {task_id} not found") + + # Analyze task complexity and determine if coordination is needed + complexity_analysis = await self._analyze_task_complexity(task_data) + + if not complexity_analysis["requires_coordination"]: + logger.info(f"Task {task_id} does not require coordination") + return None + + # Find suitable agents + if required_agents: + selected_agents = required_agents + else: + selected_agents = await self._select_agents_for_task( + task_data, required_skills, complexity_analysis + ) + + if len(selected_agents) < 2: + logger.warning( + f"Not enough agents available for coordination: {len(selected_agents)}" + ) + return None + + # Determine coordinator agent (first agent or most experienced) + coordinator_agent_id = await self._select_coordinator( + selected_agents, task_data + ) + + # Create coordination session + session_id = str(uuid.uuid4()) + session = CoordinationSession( + session_id=session_id, + root_task_id=task_id, + coordinator_agent_id=coordinator_agent_id, + participating_agents=set(selected_agents), + coordination_mode=coordination_mode, + status=CoordinationStatus.INITIALIZING, + ) + + self.active_sessions[session_id] = session + + logger.info( + f"Created coordination session {session_id} with {len(selected_agents)} agents" + ) + return session_id + + except Exception as e: + logger.error(f"Failed to initiate coordination for task {task_id}: {e}") + raise + + async def get_coordination_status( + self, session_id: str + ) -> Optional[Dict[str, Any]]: + """Get status of a coordination session""" + session = self.active_sessions.get(session_id) + if not session: + return None + + # Get subtask statuses + subtask_statuses = {} + for subtask_id, agent_id in session.subtasks.items(): + status = await self.task_execution_engine.get_execution_status(subtask_id) + subtask_statuses[subtask_id] = { + "agent_id": agent_id, + "status": status.get("status", "unknown") if status else "unknown", + } + + return { + "session_id": session_id, + "root_task_id": session.root_task_id, + "coordinator": session.coordinator_agent_id, + "participating_agents": list(session.participating_agents), + "coordination_mode": session.coordination_mode.value, + "status": session.status.value, + "subtasks": subtask_statuses, + "communications_count": len(session.communications), + "started_at": session.started_at.isoformat(), + "completed_at": ( + session.completed_at.isoformat() if session.completed_at else None + ), + "estimated_completion": ( + session.plan.estimated_completion.isoformat() if session.plan else None + ), + } + + async def send_agent_communication( + self, + from_agent_id: str, + to_agent_id: str, + message_type: str, + content: str, + metadata: Optional[Dict[str, Any]] = None, + ) -> str: + """Send communication between agents""" + + communication_id = str(uuid.uuid4()) + communication = AgentCommunication( + communication_id=communication_id, + from_agent_id=from_agent_id, + to_agent_id=to_agent_id, + message_type=message_type, + content=content, + metadata=metadata or {}, + timestamp=datetime.now(), + ) + + # Find coordination session for these agents + session = self._find_session_by_agents([from_agent_id, to_agent_id]) + if session: + session.communications.append(communication) + + # Store in database + await self._store_communication(communication) + + logger.info( + f"Agent communication {communication_id}: {from_agent_id} -> {to_agent_id}" + ) + return communication_id + + async def cancel_coordination(self, session_id: str) -> bool: + """Cancel an active coordination session""" + session = self.active_sessions.get(session_id) + if not session: + return False + + try: + # Cancel all subtasks + for subtask_id in session.subtasks.keys(): + await self.task_execution_engine.cancel_task_execution(subtask_id) + + # Update session status + session.status = CoordinationStatus.CANCELLED + session.completed_at = datetime.now() + + # Cleanup + await self._cleanup_session(session_id) + + logger.info(f"Cancelled coordination session {session_id}") + return True + + except Exception as e: + logger.error(f"Error cancelling coordination session {session_id}: {e}") + return False + + # Private methods + + async def _coordination_worker(self): + """Main coordination worker that manages session lifecycle""" + while self.running: + try: + # Process sessions that need attention + for session_id, session in list(self.active_sessions.items()): + try: + await self._process_coordination_session(session) + except Exception as e: + logger.error( + f"Error processing coordination session {session_id}: {e}" + ) + + await asyncio.sleep(5) # Check every 5 seconds + + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"Error in coordination worker: {e}") + await asyncio.sleep(10) + + async def _communication_worker(self): + """Worker that handles inter-agent communications""" + while self.running: + try: + # Process pending communications + for session in self.active_sessions.values(): + for comm in session.communications: + if not comm.response_id and comm.message_type == "request": + # Check if communication has timed out + if ( + datetime.now() - comm.timestamp + ).total_seconds() > self.communication_timeout: + await self._handle_communication_timeout(session, comm) + + await asyncio.sleep(10) # Check every 10 seconds + + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"Error in communication worker: {e}") + await asyncio.sleep(10) + + async def _synchronization_worker(self): + """Worker that synchronizes coordination sessions""" + while self.running: + try: + for session_id, session in list(self.active_sessions.items()): + if session.status == CoordinationStatus.EXECUTING: + await self._synchronize_session(session) + + await asyncio.sleep(self.synchronization_interval) + + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"Error in synchronization worker: {e}") + await asyncio.sleep(self.synchronization_interval) + + async def _process_coordination_session(self, session: CoordinationSession): + """Process a coordination session based on its current status""" + + if session.status == CoordinationStatus.INITIALIZING: + await self._initialize_session(session) + elif session.status == CoordinationStatus.PLANNING: + await self._plan_coordination(session) + elif session.status == CoordinationStatus.EXECUTING: + await self._monitor_execution(session) + elif session.status == CoordinationStatus.REVIEWING: + await self._review_coordination(session) + + async def _initialize_session(self, session: CoordinationSession): + """Initialize a coordination session""" + try: + # Get detailed task information + task_data = await DatabaseManager.get_task(session.root_task_id) + + # Update agent capabilities + await self._update_agent_capabilities(list(session.participating_agents)) + + # Move to planning phase + session.status = CoordinationStatus.PLANNING + + logger.info(f"Initialized coordination session {session.session_id}") + + except Exception as e: + logger.error(f"Error initializing session {session.session_id}: {e}") + session.status = CoordinationStatus.FAILED + + async def _plan_coordination(self, session: CoordinationSession): + """Create coordination plan""" + try: + # Get task data + task_data = await DatabaseManager.get_task(session.root_task_id) + + # Decompose task into subtasks + subtasks = await self._decompose_task(task_data, session.coordination_mode) + + # Assign agents to subtasks + assignments = await self._assign_agents_to_subtasks( + subtasks, list(session.participating_agents) + ) + + # Create dependencies + dependencies = await self._create_task_dependencies( + subtasks, session.coordination_mode + ) + + # Estimate completion time + estimated_completion = await self._estimate_coordination_completion( + subtasks, assignments, dependencies + ) + + # Create coordination plan + plan = CoordinationPlan( + plan_id=str(uuid.uuid4()), + root_task_id=session.root_task_id, + coordination_mode=session.coordination_mode, + participating_agents=list(session.participating_agents), + task_assignments=assignments, + dependencies=dependencies, + estimated_completion=estimated_completion, + created_at=datetime.now(), + ) + + session.plan = plan + session.status = CoordinationStatus.EXECUTING + + # Create subtasks in database and start execution + for subtask_data in subtasks: + subtask_id = await self._create_subtask(subtask_data, assignments) + session.subtasks[subtask_id] = assignments[subtask_data["id"]] + + # Start subtask execution + await self.task_execution_engine.start_task_execution(subtask_id) + + logger.info( + f"Created coordination plan for session {session.session_id} with {len(subtasks)} subtasks" + ) + + except Exception as e: + logger.error(f"Error planning coordination {session.session_id}: {e}") + session.status = CoordinationStatus.FAILED + + async def _monitor_execution(self, session: CoordinationSession): + """Monitor execution of coordinated tasks""" + try: + # Check status of all subtasks + completed_subtasks = 0 + failed_subtasks = 0 + + for subtask_id in session.subtasks.keys(): + status = await self.task_execution_engine.get_execution_status( + subtask_id + ) + if status: + if status.get("status") == "completed": + completed_subtasks += 1 + elif status.get("status") == "failed": + failed_subtasks += 1 + + total_subtasks = len(session.subtasks) + + # Check if coordination is complete + if completed_subtasks == total_subtasks: + session.status = CoordinationStatus.REVIEWING + elif failed_subtasks > 0: + # Handle failures + await self._handle_coordination_failures(session) + + except Exception as e: + logger.error(f"Error monitoring execution {session.session_id}: {e}") + + async def _review_coordination(self, session: CoordinationSession): + """Review completed coordination and aggregate results""" + try: + # Collect results from all subtasks + results = {} + for subtask_id, agent_id in session.subtasks.items(): + status = await self.task_execution_engine.get_execution_status( + subtask_id + ) + if status: + results[subtask_id] = { + "agent_id": agent_id, + "status": status.get("status"), + "result": status.get("result"), + } + + # Aggregate results + coordination_result = await self._aggregate_coordination_results(results) + + # Complete coordination + session.status = CoordinationStatus.COMPLETED + session.completed_at = datetime.now() + session.result = coordination_result + + # Update root task status + await DatabaseManager.update_task_status( + session.root_task_id, "completed", coordination_result + ) + + logger.info(f"Completed coordination session {session.session_id}") + + # Schedule cleanup + asyncio.create_task(self._cleanup_session(session.session_id)) + + except Exception as e: + logger.error(f"Error reviewing coordination {session.session_id}: {e}") + session.status = CoordinationStatus.FAILED + + async def _analyze_task_complexity( + self, task_data: Dict[str, Any] + ) -> Dict[str, Any]: + """Analyze task complexity to determine if coordination is needed""" + + description = task_data.get("description", "") + title = task_data.get("title", "") + + # Simple heuristics for complexity analysis + complexity_indicators = [ + "multiple components" in description.lower(), + "frontend and backend" in description.lower(), + "database and api" in description.lower(), + "testing and deployment" in description.lower(), + len(description.split()) > 50, # Long description + "integrate" in description.lower(), + "coordinate" in description.lower(), + "collaborate" in description.lower(), + ] + + complexity_score = sum(complexity_indicators) / len(complexity_indicators) + + return { + "requires_coordination": complexity_score > 0.3, + "complexity_score": complexity_score, + "estimated_agents_needed": min(max(2, int(complexity_score * 5)), 5), + "estimated_duration_hours": max(4, int(complexity_score * 24)), + } + + async def _select_agents_for_task( + self, + task_data: Dict[str, Any], + required_skills: Optional[List[str]], + complexity_analysis: Dict[str, Any], + ) -> List[str]: + """Select appropriate agents for the task""" + + # Get all available agents + all_agents = await DatabaseManager.get_all_agents() + + # Filter by availability and skills + suitable_agents = [] + for agent in all_agents: + if agent["status"] == "available": + agent_skills = agent.get("config", {}).get("tools", []) + + # Check skill match + if required_skills: + skill_match = any( + skill in agent_skills for skill in required_skills + ) + else: + skill_match = True + + if skill_match: + suitable_agents.append(agent["id"]) + + # Select optimal number of agents + max_agents = complexity_analysis.get("estimated_agents_needed", 3) + return suitable_agents[:max_agents] + + async def _select_coordinator( + self, agents: List[str], task_data: Dict[str, Any] + ) -> str: + """Select the coordinator agent from available agents""" + + # For now, select the first agent as coordinator + # In production, this would consider agent experience, current load, etc. + return agents[0] + + async def _decompose_task( + self, task_data: Dict[str, Any], mode: CoordinationMode + ) -> List[Dict[str, Any]]: + """Decompose a complex task into subtasks""" + + # Simple task decomposition based on common patterns + description = task_data.get("description", "") + title = task_data.get("title", "") + + subtasks = [] + + # Common subtask patterns + if "frontend" in description.lower() or "ui" in description.lower(): + subtasks.append( + { + "id": f"frontend-{uuid.uuid4()}", + "title": f"Frontend Implementation - {title}", + "description": f"Implement frontend components for: {description}", + "type": "frontend_development", + "estimated_hours": 4, + } + ) + + if "backend" in description.lower() or "api" in description.lower(): + subtasks.append( + { + "id": f"backend-{uuid.uuid4()}", + "title": f"Backend Implementation - {title}", + "description": f"Implement backend services for: {description}", + "type": "backend_development", + "estimated_hours": 6, + } + ) + + if "database" in description.lower() or "data" in description.lower(): + subtasks.append( + { + "id": f"database-{uuid.uuid4()}", + "title": f"Database Design - {title}", + "description": f"Design and implement database schema for: {description}", + "type": "database_development", + "estimated_hours": 3, + } + ) + + if "test" in description.lower(): + subtasks.append( + { + "id": f"testing-{uuid.uuid4()}", + "title": f"Testing - {title}", + "description": f"Create comprehensive tests for: {description}", + "type": "testing", + "estimated_hours": 4, + } + ) + + # If no specific subtasks identified, create generic subtasks + if not subtasks: + subtasks = [ + { + "id": f"analysis-{uuid.uuid4()}", + "title": f"Analysis - {title}", + "description": f"Analyze requirements for: {description}", + "type": "analysis", + "estimated_hours": 2, + }, + { + "id": f"implementation-{uuid.uuid4()}", + "title": f"Implementation - {title}", + "description": f"Implement solution for: {description}", + "type": "implementation", + "estimated_hours": 6, + }, + { + "id": f"review-{uuid.uuid4()}", + "title": f"Review - {title}", + "description": f"Review and validate solution for: {description}", + "type": "review", + "estimated_hours": 2, + }, + ] + + return subtasks + + async def _assign_agents_to_subtasks( + self, subtasks: List[Dict[str, Any]], agents: List[str] + ) -> Dict[str, str]: + """Assign agents to subtasks based on capabilities""" + + assignments = {} + + # Get agent capabilities + agent_data = {} + for agent_id in agents: + agent = await DatabaseManager.get_agent(agent_id) + if agent: + agent_data[agent_id] = agent + + # Simple assignment based on agent type + for subtask in subtasks: + subtask_type = subtask.get("type", "") + best_agent = None + + # Match agent type to subtask type + for agent_id, agent in agent_data.items(): + agent_type = agent.get("type", "") + + if subtask_type.startswith("frontend") and "frontend" in agent_type: + best_agent = agent_id + break + elif subtask_type.startswith("backend") and "backend" in agent_type: + best_agent = agent_id + break + elif subtask_type.startswith("database") and "backend" in agent_type: + best_agent = agent_id + break + elif subtask_type == "testing" and "qa" in agent_type: + best_agent = agent_id + break + + # Fallback to first available agent + if not best_agent: + best_agent = agents[0] + + assignments[subtask["id"]] = best_agent + + return assignments + + async def _create_task_dependencies( + self, subtasks: List[Dict[str, Any]], mode: CoordinationMode + ) -> List[TaskDependency]: + """Create dependencies between subtasks""" + + dependencies = [] + + if mode == CoordinationMode.SEQUENTIAL: + # Create sequential dependencies + for i in range(1, len(subtasks)): + dependencies.append( + TaskDependency( + dependent_task_id=subtasks[i]["id"], + prerequisite_task_id=subtasks[i - 1]["id"], + dependency_type="blocking", + ) + ) + + elif mode == CoordinationMode.HIERARCHICAL: + # Analysis task should complete before implementation + analysis_tasks = [t for t in subtasks if "analysis" in t["type"]] + implementation_tasks = [ + t for t in subtasks if "implementation" in t["type"] + ] + + for impl_task in implementation_tasks: + for analysis_task in analysis_tasks: + dependencies.append( + TaskDependency( + dependent_task_id=impl_task["id"], + prerequisite_task_id=analysis_task["id"], + dependency_type="blocking", + ) + ) + + # PARALLEL and COLLABORATIVE modes have no strict dependencies + + return dependencies + + async def _estimate_coordination_completion( + self, + subtasks: List[Dict[str, Any]], + assignments: Dict[str, str], + dependencies: List[TaskDependency], + ) -> datetime: + """Estimate when coordination will complete""" + + if not dependencies: + # Parallel execution - completion time is max of all subtasks + max_hours = max(subtask.get("estimated_hours", 4) for subtask in subtasks) + else: + # Sequential/dependent execution - sum of critical path + total_hours = sum(subtask.get("estimated_hours", 4) for subtask in subtasks) + max_hours = min(total_hours, 24) # Cap at 24 hours + + return datetime.now() + timedelta(hours=max_hours) + + async def _create_subtask( + self, subtask_data: Dict[str, Any], assignments: Dict[str, str] + ) -> str: + """Create a subtask in the database""" + + subtask_id = str(uuid.uuid4()) + agent_id = assignments.get(subtask_data["id"]) + + # Create task in database + await DatabaseManager.create_task( + task_id=subtask_id, + title=subtask_data["title"], + description=subtask_data["description"], + assigned_to=agent_id, + priority="medium", + metadata={ + "coordination_subtask": True, + "parent_task_type": subtask_data.get("type"), + "estimated_hours": subtask_data.get("estimated_hours", 4), + }, + ) + + return subtask_id + + def _find_session_by_agents( + self, agent_ids: List[str] + ) -> Optional[CoordinationSession]: + """Find coordination session that includes the specified agents""" + for session in self.active_sessions.values(): + if any(agent_id in session.participating_agents for agent_id in agent_ids): + return session + return None + + async def _store_communication(self, communication: AgentCommunication): + """Store agent communication in database""" + # This would store the communication in the database + # For now, just log it + logger.info( + f"Agent communication: {communication.from_agent_id} -> {communication.to_agent_id}: {communication.content}" + ) + + async def _handle_communication_timeout( + self, session: CoordinationSession, communication: AgentCommunication + ): + """Handle communication timeout""" + logger.warning( + f"Communication timeout in session {session.session_id}: {communication.communication_id}" + ) + + # Could implement retry logic or escalation here + + async def _synchronize_session(self, session: CoordinationSession): + """Synchronize coordination session state""" + # Check if any agents need help or coordination + # Update session status based on subtask progress + # Handle any conflicts or issues + pass + + async def _handle_coordination_failures(self, session: CoordinationSession): + """Handle failures in coordination""" + logger.warning( + f"Handling failures in coordination session {session.session_id}" + ) + + # Could implement retry logic, reassignment, or escalation + session.status = CoordinationStatus.FAILED + + async def _aggregate_coordination_results( + self, results: Dict[str, Dict[str, Any]] + ) -> Dict[str, Any]: + """Aggregate results from all coordinated subtasks""" + + successful_subtasks = [ + r for r in results.values() if r["status"] == "completed" + ] + + return { + "coordination_completed": True, + "total_subtasks": len(results), + "successful_subtasks": len(successful_subtasks), + "failed_subtasks": len(results) - len(successful_subtasks), + "results": results, + "completion_time": datetime.now().isoformat(), + } + + async def _update_agent_capabilities(self, agent_ids: List[str]): + """Update cached agent capabilities""" + for agent_id in agent_ids: + agent_data = await DatabaseManager.get_agent(agent_id) + if agent_data: + # Extract capabilities from agent configuration + tools = agent_data.get("config", {}).get("tools", []) + capabilities = [ + AgentCapability( + skill=tool, + proficiency=0.8, # Default proficiency + availability=agent_data.get("status") == "available", + current_load=0.5, # Default load + ) + for tool in tools + ] + self.agent_capabilities[agent_id] = capabilities + + async def _cleanup_session(self, session_id: str): + """Clean up completed coordination session""" + session = self.active_sessions.pop(session_id, None) + if session: + logger.info(f"Cleaned up coordination session {session_id}") + + +# Integration with existing TaskExecutionEngine +def integrate_multi_agent_coordination( + task_execution_engine: TaskExecutionEngine, +) -> MultiAgentCoordinator: + """Create and integrate multi-agent coordinator with task execution engine""" + coordinator = MultiAgentCoordinator(task_execution_engine) + + # Add coordination capabilities to task execution engine + task_execution_engine.multi_agent_coordinator = coordinator + + return coordinator diff --git a/services/orchestrator/task_execution_engine.py b/services/orchestrator/task_execution_engine.py index b2f3b3a..b11cd75 100644 --- a/services/orchestrator/task_execution_engine.py +++ b/services/orchestrator/task_execution_engine.py @@ -1,1306 +1,1306 @@ -""" -Task Execution Engine for FuzeAgent Autonomous Execution - -Orchestrates the autonomous execution of tasks by agents, managing: -- Task lifecycle and state transitions -- Sandbox creation and cleanup -- Git workflow automation -- Human-in-the-loop interactions -- Inter-agent communication -- Result aggregation - -This is the core component that ties together all autonomous execution components. -""" - -import asyncio -import json -import logging -import uuid -from dataclasses import dataclass -from datetime import datetime, timedelta -from enum import Enum -from typing import Any, Callable, Dict, List, Optional - -from .claude_code_wrapper import ClaudeCodeWrapper -from .claude_sdk_manager import ClaudeSDKManager, ClaudeSDKSession -from .context_enhancement_service import ContextEnhancementService -from .conversation_manager import ConversationManager, InteractionType -from .database import DatabaseManager, get_db_connection -from .file_operations_engine import FileOperationsEngine -from .git_workflow_manager import GitWorkflowManager -from .sandbox_manager import AgentSandboxManager, Sandbox -from .task_knowledge_extractor import TaskKnowledgeExtractor - -logger = logging.getLogger(__name__) - - -class TaskStatus(str, Enum): - PENDING = "pending" - ANALYZING = "analyzing" - SETTING_UP = "setting_up" - EXECUTING = "executing" - WAITING_FOR_HUMAN = "waiting_for_human" - REVIEWING = "reviewing" - COMMITTING = "committing" - COMPLETED = "completed" - FAILED = "failed" - CANCELLED = "cancelled" - - -class ExecutionStep(str, Enum): - ANALYZE_TASK = "analyze_task" - SETUP_SANDBOX = "setup_sandbox" - SETUP_GIT = "setup_git" - EXECUTE_ITERATION = "execute_iteration" - REVIEW_CHANGES = "review_changes" - COMMIT_CHANGES = "commit_changes" - HUMAN_INTERACTION = "human_interaction" - FINALIZE_TASK = "finalize_task" - CLEANUP = "cleanup" - - -@dataclass -class TaskIteration: - """Represents a single iteration of task execution""" - - iteration_number: int - step: ExecutionStep - started_at: datetime - completed_at: Optional[datetime] - input_data: Dict[str, Any] - output_data: Optional[Dict[str, Any]] - success: bool - error_message: Optional[str] - human_question: Optional[str] = None - human_response: Optional[str] = None - - -@dataclass -class ExecutionContext: - """Context for task execution""" - - task_id: str - agent_id: str - task_data: Dict[str, Any] - agent_data: Dict[str, Any] - sandbox: Optional[Sandbox] - git_manager: Optional[GitWorkflowManager] - claude_wrapper: Optional[ClaudeCodeWrapper] - current_iteration: int - iterations: List[TaskIteration] - status: TaskStatus - started_at: datetime - completed_at: Optional[datetime] - result: Optional[Dict[str, Any]] - error: Optional[str] - # New components for autonomous execution - file_operations_engine: Optional[FileOperationsEngine] = None - claude_sdk_manager: Optional[ClaudeSDKManager] = None - claude_session_id: Optional[str] = None - - -class TaskExecutionEngine: - """ - Orchestrates autonomous task execution by agents. - - Features: - - Task lifecycle management - - Sandbox and Git workflow integration - - Human-in-the-loop interactions - - Dependency handling - - Result aggregation - - Error recovery - """ - - def __init__( - self, - sandbox_manager: AgentSandboxManager, - knowledge_extractor: Optional[TaskKnowledgeExtractor] = None, - context_enhancer: Optional[ContextEnhancementService] = None, - ): - self.sandbox_manager = sandbox_manager - self.conversation_manager = ConversationManager() - self.active_executions: Dict[str, ExecutionContext] = {} - self.execution_callbacks: Dict[str, List[Callable]] = {} - self.running = False - self.worker_tasks: List[asyncio.Task] = [] - - # Knowledge management services - self.knowledge_extractor = knowledge_extractor - self.context_enhancer = context_enhancer - - # Initialize integrated components - self.file_operations_engines: Dict[str, FileOperationsEngine] = {} # Per task - self.claude_sdk_managers: Dict[str, ClaudeSDKManager] = {} # Per task - - # Configuration - self.max_iterations = 50 - self.iteration_timeout = 3600 # 1 hour per iteration - self.human_response_timeout = 86400 # 24 hours for human response - - async def start(self): - """Start the execution engine""" - logger.info("Starting TaskExecutionEngine") - self.running = True - - # Start worker tasks - self.worker_tasks = [ - asyncio.create_task(self._execution_worker()), - asyncio.create_task(self._monitoring_worker()), - asyncio.create_task(self._cleanup_worker()), - ] - - logger.info("TaskExecutionEngine started") - - async def stop(self): - """Stop the execution engine""" - logger.info("Stopping TaskExecutionEngine") - self.running = False - - # Cancel worker tasks - for task in self.worker_tasks: - task.cancel() - - try: - await asyncio.gather(*self.worker_tasks, return_exceptions=True) - except Exception as e: - logger.error(f"Error stopping worker tasks: {e}") - - # Clean up active executions - for execution_id in list(self.active_executions.keys()): - try: - await self._cleanup_execution(execution_id) - except Exception as e: - logger.error(f"Error cleaning up execution {execution_id}: {e}") - - logger.info("TaskExecutionEngine stopped") - - async def start_task_execution(self, task_id: str) -> Dict[str, Any]: - """ - Start autonomous execution of a task. - Returns execution status and context. - """ - logger.info(f"Starting task execution: {task_id}") - - try: - # Get task data - task_data = await self._get_task_data(task_id) - if not task_data: - raise ValueError(f"Task {task_id} not found") - - # Get agent data - agent_id = task_data.get("assigned_to") - if not agent_id: - raise ValueError(f"Task {task_id} has no assigned agent") - - agent_data = await self._get_agent_data(agent_id) - if not agent_data: - raise ValueError(f"Agent {agent_id} not found") - - # Create execution context - execution_context = ExecutionContext( - task_id=task_id, - agent_id=agent_id, - task_data=task_data, - agent_data=agent_data, - sandbox=None, - git_manager=None, - claude_wrapper=None, - current_iteration=0, - iterations=[], - status=TaskStatus.PENDING, - started_at=datetime.now(), - completed_at=None, - result=None, - error=None, - ) - - # Store execution context - self.active_executions[task_id] = execution_context - - # Update task status in database - await self._update_task_status(task_id, TaskStatus.PENDING) - - logger.info(f"✅ Task execution started: {task_id}") - return { - "task_id": task_id, - "status": TaskStatus.PENDING.value, - "execution_started": True, - "agent_id": agent_id, - } - - except Exception as e: - logger.error(f"❌ Failed to start task execution {task_id}: {e}") - await self._update_task_status(task_id, TaskStatus.FAILED, error=str(e)) - raise - - async def get_execution_status(self, task_id: str) -> Dict[str, Any]: - """Get detailed execution status for a task""" - - execution = self.active_executions.get(task_id) - if not execution: - # Check database for completed/failed tasks - task_data = await self._get_task_data(task_id) - if task_data: - return { - "task_id": task_id, - "status": task_data.get("status", "unknown"), - "result": task_data.get("result"), - "active_execution": False, - } - else: - return {"task_id": task_id, "status": "not_found"} - - return { - "task_id": task_id, - "status": execution.status.value, - "agent_id": execution.agent_id, - "current_iteration": execution.current_iteration, - "iterations_count": len(execution.iterations), - "started_at": execution.started_at.isoformat(), - "completed_at": ( - execution.completed_at.isoformat() if execution.completed_at else None - ), - "sandbox_id": execution.sandbox.sandbox_id if execution.sandbox else None, - "git_branch": ( - execution.git_manager.feature_branch if execution.git_manager else None - ), - "result": execution.result, - "error": execution.error, - "active_execution": True, - } - - async def get_task_iterations(self, task_id: str) -> List[Dict[str, Any]]: - """Get iteration history for a task""" - - execution = self.active_executions.get(task_id) - if execution: - iterations = execution.iterations - else: - # Get from database - iterations = await self._get_task_iterations_from_db(task_id) - - return [ - { - "iteration_number": it.iteration_number, - "step": it.step.value if hasattr(it.step, "value") else str(it.step), - "started_at": it.started_at.isoformat(), - "completed_at": ( - it.completed_at.isoformat() if it.completed_at else None - ), - "success": it.success, - "error_message": it.error_message, - "human_question": it.human_question, - "human_response": it.human_response, - "input_data": it.input_data, - "output_data": it.output_data, - } - for it in iterations - ] - - async def cancel_task_execution(self, task_id: str) -> bool: - """Cancel a running task execution""" - - execution = self.active_executions.get(task_id) - if not execution: - return False - - execution.status = TaskStatus.CANCELLED - execution.completed_at = datetime.now() - execution.error = "Task cancelled by user" - - # Update database - await self._update_task_status( - task_id, TaskStatus.CANCELLED, error="Task cancelled by user" - ) - - # Schedule cleanup - asyncio.create_task(self._cleanup_execution(task_id)) - - logger.info(f"Task execution cancelled: {task_id}") - return True - - # Private methods for execution workflow - - async def _execution_worker(self): - """Main execution worker that processes pending tasks""" - while self.running: - try: - # Find tasks ready for execution - pending_tasks = [ - task_id - for task_id, execution in self.active_executions.items() - if execution.status in [TaskStatus.PENDING, TaskStatus.EXECUTING] - ] - - # Process each pending task - for task_id in pending_tasks: - try: - await self._process_task_execution(task_id) - except Exception as e: - logger.error(f"Error processing task {task_id}: {e}") - await self._handle_execution_error(task_id, str(e)) - - # Sleep between iterations - await asyncio.sleep(5) - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Error in execution worker: {e}") - await asyncio.sleep(10) - - async def _process_task_execution(self, task_id: str): - """Process a single task execution step""" - execution = self.active_executions.get(task_id) - if not execution: - return - - # Skip if waiting for human or in terminal state - if execution.status in [ - TaskStatus.WAITING_FOR_HUMAN, - TaskStatus.COMPLETED, - TaskStatus.FAILED, - TaskStatus.CANCELLED, - ]: - return - - # Determine next step - next_step = self._determine_next_step(execution) - if not next_step: - return - - # Execute the step - try: - await self._execute_step(execution, next_step) - except Exception as e: - logger.error(f"Error executing step {next_step} for task {task_id}: {e}") - await self._handle_execution_error(task_id, str(e)) - - def _determine_next_step( - self, execution: ExecutionContext - ) -> Optional[ExecutionStep]: - """Determine the next execution step""" - - if execution.status == TaskStatus.PENDING: - return ExecutionStep.ANALYZE_TASK - - if not execution.iterations: - return ExecutionStep.ANALYZE_TASK - - last_iteration = execution.iterations[-1] - - # Continue based on last completed step - if last_iteration.step == ExecutionStep.ANALYZE_TASK and last_iteration.success: - return ExecutionStep.SETUP_SANDBOX - elif ( - last_iteration.step == ExecutionStep.SETUP_SANDBOX - and last_iteration.success - ): - return ExecutionStep.SETUP_GIT - elif last_iteration.step == ExecutionStep.SETUP_GIT and last_iteration.success: - return ExecutionStep.EXECUTE_ITERATION - elif ( - last_iteration.step == ExecutionStep.EXECUTE_ITERATION - and last_iteration.success - ): - # Check if we need human input - if last_iteration.human_question: - return ExecutionStep.HUMAN_INTERACTION - else: - return ExecutionStep.REVIEW_CHANGES - elif ( - last_iteration.step == ExecutionStep.HUMAN_INTERACTION - and last_iteration.human_response - ): - return ExecutionStep.EXECUTE_ITERATION - elif ( - last_iteration.step == ExecutionStep.REVIEW_CHANGES - and last_iteration.success - ): - return ExecutionStep.COMMIT_CHANGES - elif ( - last_iteration.step == ExecutionStep.COMMIT_CHANGES - and last_iteration.success - ): - # Check if task is complete - if self._is_task_complete(execution): - return ExecutionStep.FINALIZE_TASK - else: - return ExecutionStep.EXECUTE_ITERATION - - return None - - async def _execute_step(self, execution: ExecutionContext, step: ExecutionStep): - """Execute a specific step""" - - iteration = TaskIteration( - iteration_number=execution.current_iteration + 1, - step=step, - started_at=datetime.now(), - completed_at=None, - input_data={}, - output_data=None, - success=False, - error_message=None, - ) - - execution.iterations.append(iteration) - execution.current_iteration += 1 - - try: - if step == ExecutionStep.ANALYZE_TASK: - await self._step_analyze_task(execution, iteration) - elif step == ExecutionStep.SETUP_SANDBOX: - await self._step_setup_sandbox(execution, iteration) - elif step == ExecutionStep.SETUP_GIT: - await self._step_setup_git(execution, iteration) - elif step == ExecutionStep.EXECUTE_ITERATION: - await self._step_execute_iteration(execution, iteration) - elif step == ExecutionStep.REVIEW_CHANGES: - await self._step_review_changes(execution, iteration) - elif step == ExecutionStep.COMMIT_CHANGES: - await self._step_commit_changes(execution, iteration) - elif step == ExecutionStep.HUMAN_INTERACTION: - await self._step_human_interaction(execution, iteration) - elif step == ExecutionStep.FINALIZE_TASK: - await self._step_finalize_task(execution, iteration) - - iteration.completed_at = datetime.now() - iteration.success = True - - except Exception as e: - iteration.completed_at = datetime.now() - iteration.success = False - iteration.error_message = str(e) - raise - - finally: - # Store iteration in database - await self._store_task_iteration(execution.task_id, iteration) - - async def _step_analyze_task( - self, execution: ExecutionContext, iteration: TaskIteration - ): - """Analyze the task and prepare execution plan""" - execution.status = TaskStatus.ANALYZING - await self._update_task_status(execution.task_id, TaskStatus.ANALYZING) - - # Analyze task requirements - task_description = execution.task_data.get("description", "") - task_title = execution.task_data.get("title", "") - - iteration.input_data = { - "task_title": task_title, - "task_description": task_description, - "agent_type": execution.agent_data.get("type"), - "agent_role": execution.agent_data.get("role"), - } - - # Enhance context with organizational knowledge - enhanced_context = None - if self.context_enhancer: - try: - enhanced_context = await self.context_enhancer.enhance_agent_context( - agent_id=execution.agent_id, - task_data=execution.task_data, - base_context=iteration.input_data, - ) - logger.info( - f"Enhanced context for task {execution.task_id}: " - f"{len(enhanced_context.organizational_knowledge)} org + " - f"{len(enhanced_context.team_knowledge)} team + " - f"{len(enhanced_context.similar_task_insights)} similar task insights" - ) - except Exception as e: - logger.error( - f"Error enhancing context for task {execution.task_id}: {e}" - ) - - # Simple analysis for now - in a full implementation this would use AI - iteration.output_data = { - "analysis_complete": True, - "requires_sandbox": execution.agent_data.get("type") == "developer", - "requires_git": bool( - execution.agent_data.get("repository_settings", {}).get( - "repository_url" - ) - ), - "enhanced_context": enhanced_context, - "estimated_complexity": "medium", - "estimated_iterations": 5, - } - - logger.info(f"Task analysis complete for {execution.task_id}") - - async def _step_setup_sandbox( - self, execution: ExecutionContext, iteration: TaskIteration - ): - """Set up sandbox environment for the agent""" - execution.status = TaskStatus.SETTING_UP - await self._update_task_status(execution.task_id, TaskStatus.SETTING_UP) - - agent_template = execution.agent_data.get("template_id", "python_developer") - repository_settings = execution.agent_data.get("repository_settings", {}) - sandbox_settings = execution.agent_data.get("sandbox_settings", {}) - - # Create sandbox - sandbox = await self.sandbox_manager.create_sandbox( - agent_id=execution.agent_id, - task_id=execution.task_id, - agent_template=agent_template, - repository_settings=repository_settings, - custom_settings=sandbox_settings, - ) - - execution.sandbox = sandbox - - iteration.input_data = { - "agent_template": agent_template, - "repository_settings": repository_settings, - "sandbox_settings": sandbox_settings, - } - - iteration.output_data = { - "sandbox_id": sandbox.sandbox_id, - "workspace_path": sandbox.workspace_path, - "container_id": sandbox.container_id, - } - - logger.info( - f"Sandbox setup complete for {execution.task_id}: {sandbox.sandbox_id}" - ) - - async def _step_setup_git( - self, execution: ExecutionContext, iteration: TaskIteration - ): - """Set up Git workflow for the task""" - repository_settings = execution.agent_data.get("repository_settings", {}) - - if not repository_settings.get("repository_url"): - # Skip Git setup if no repository - iteration.output_data = { - "git_setup": "skipped", - "reason": "no_repository_configured", - } - return - - # Create Git workflow manager - git_manager = GitWorkflowManager( - agent_id=execution.agent_id, - task_id=execution.task_id, - repo_settings=repository_settings, - ) - - # Setup workspace - feature_branch = await git_manager.setup_workspace() - - execution.git_manager = git_manager - - # Create enhanced Claude wrapper with Git context and conversation tracking - execution.claude_wrapper = ClaudeCodeWrapper( - workspace_path=git_manager.workspace_path, - git_manager=git_manager, - agent_id=execution.agent_id, - task_id=execution.task_id, - conversation_manager=self.conversation_manager, - ) - - # Initialize File Operations Engine - file_ops_engine = FileOperationsEngine(git_manager.workspace_path) - execution.file_operations_engine = file_ops_engine - self.file_operations_engines[execution.task_id] = file_ops_engine - - # Initialize Claude SDK Manager - claude_sdk_manager = ClaudeSDKManager( - file_operations_engine=file_ops_engine, - conversation_manager=self.conversation_manager, - ) - execution.claude_sdk_manager = claude_sdk_manager - self.claude_sdk_managers[execution.task_id] = claude_sdk_manager - - # Start conversation session - await execution.claude_wrapper.start_conversation_session( - execution.sandbox.sandbox_id - ) - - iteration.input_data = { - "repository_url": repository_settings.get("repository_url"), - "default_branch": repository_settings.get("default_branch", "main"), - } - - iteration.output_data = { - "git_setup": "complete", - "feature_branch": feature_branch, - "workspace_path": git_manager.workspace_path, - } - - logger.info(f"Git setup complete for {execution.task_id}: {feature_branch}") - - async def _step_execute_iteration( - self, execution: ExecutionContext, iteration: TaskIteration - ): - """Execute a development iteration using Claude SDK Manager""" - execution.status = TaskStatus.EXECUTING - await self._update_task_status(execution.task_id, TaskStatus.EXECUTING) - - task_description = execution.task_data.get("description", "") - task_title = execution.task_data.get("title", "") - - iteration.input_data = { - "task_description": task_description, - "task_title": task_title, - "iteration_number": iteration.iteration_number, - "workspace_path": ( - execution.sandbox.workspace_path if execution.sandbox else None - ), - } - - try: - # Start Claude SDK session if not already running - if not execution.claude_session_id and execution.claude_sdk_manager: - # Build comprehensive task context - context_info = "" - if execution.git_manager: - context_info += f"\nRepository: {execution.git_manager.repo_url}" - context_info += f"\nBranch: {execution.git_manager.feature_branch}" - if iteration.iteration_number > 1: - context_info += f"\nIteration: {iteration.iteration_number} of ongoing development" - - # Construct task prompt - task_prompt = f""" -Task: {task_title} - -Description: {task_description} - -Context: {context_info} - -Please analyze the codebase, understand the requirements, and implement the necessary changes. -Work incrementally and ask for clarification if needed. -""" - - # Start Claude SDK session - execution.claude_session_id = ( - await execution.claude_sdk_manager.start_session( - task_id=execution.task_id, - agent_id=execution.agent_id, - workspace_path=( - execution.git_manager.workspace_path - if execution.git_manager - else execution.sandbox.workspace_path - ), - task_description=task_prompt, - additional_context=context_info, - ) - ) - - logger.info( - f"Started Claude SDK session: {execution.claude_session_id}" - ) - - # Register interaction callback for human-in-the-loop - if execution.claude_sdk_manager and execution.claude_session_id: - execution.claude_sdk_manager.register_interaction_callback( - execution.claude_session_id, - lambda session, interaction: self._handle_claude_interaction( - execution, session, interaction - ), - ) - - # Monitor session status - session_status = ( - await execution.claude_sdk_manager.get_session_status( - execution.claude_session_id - ) - if execution.claude_session_id - else None - ) - - if session_status: - current_interaction = session_status.get("current_interaction") - if current_interaction: - # Claude is waiting for human input - interaction_type = current_interaction.get("type") - if interaction_type in ["user_input", "confirmation"]: - iteration.human_question = current_interaction.get("prompt") - execution.status = TaskStatus.WAITING_FOR_HUMAN - await self._update_task_status( - execution.task_id, TaskStatus.WAITING_FOR_HUMAN - ) - - iteration.output_data = { - "claude_session_status": session_status.get("state"), - "human_interaction_required": True, - "interaction_type": interaction_type, - "human_question_asked": True, - } - - logger.info( - f"Claude SDK requesting human input for task {execution.task_id}" - ) - return - - elif interaction_type == "file_approval": - # File operations pending approval - batch_id = current_interaction.get("metadata", {}).get( - "batch_id" - ) - if batch_id and execution.file_operations_engine: - # Get diff preview for human review - diffs = await execution.file_operations_engine.get_file_diff_preview( - batch_id - ) - - iteration.human_question = f"""Claude wants to make the following file changes: - -{current_interaction.get('prompt')} - -File changes preview: -{self._format_diffs_for_human(diffs)} - -Approve these changes? (yes/no)""" - - execution.status = TaskStatus.WAITING_FOR_HUMAN - await self._update_task_status( - execution.task_id, TaskStatus.WAITING_FOR_HUMAN - ) - - iteration.output_data = { - "claude_session_status": session_status.get("state"), - "file_approval_required": True, - "batch_id": batch_id, - "file_changes_preview": diffs, - "human_question_asked": True, - } - - logger.info( - f"Claude SDK requesting file approval for task {execution.task_id}" - ) - return - - # No interaction needed - continue execution - iteration.output_data = { - "claude_session_status": session_status.get("state"), - "session_active": True, - "iteration_completed": True, - "workspace_path": session_status.get("workspace_path"), - } - - # Check if session completed - if session_status.get("state") in ["completed", "terminated"]: - iteration.output_data["development_complete"] = True - - else: - # No session - this shouldn't happen but handle gracefully - iteration.output_data = { - "error": "No Claude SDK session available", - "development_complete": False, - } - - except Exception as e: - logger.error(f"Error in Claude SDK iteration: {e}") - iteration.output_data = {"error": str(e), "development_complete": False} - raise - - logger.info( - f"Development iteration {iteration.iteration_number} processed for {execution.task_id}" - ) - - async def _step_review_changes( - self, execution: ExecutionContext, iteration: TaskIteration - ): - """Review the changes made in the iteration""" - execution.status = TaskStatus.REVIEWING - await self._update_task_status(execution.task_id, TaskStatus.REVIEWING) - - # Review changes - in full implementation this would: - # 1. Run linting and type checking - # 2. Run tests - # 3. Check code quality - # 4. Validate against requirements - - iteration.output_data = { - "review_passed": True, - "issues_found": [], - "tests_passed": True, - "code_quality_score": 85, - } - - logger.info(f"Code review complete for {execution.task_id}") - - async def _step_commit_changes( - self, execution: ExecutionContext, iteration: TaskIteration - ): - """Commit changes to Git""" - execution.status = TaskStatus.COMMITTING - await self._update_task_status(execution.task_id, TaskStatus.COMMITTING) - - if not execution.git_manager: - iteration.output_data = {"commit": "skipped", "reason": "no_git_manager"} - return - - # Commit changes - commit_message = f"Iteration {iteration.iteration_number}: {execution.task_data.get('title', 'Task update')}" - commit_hash = await execution.git_manager.commit_changes( - message=commit_message, iteration_number=iteration.iteration_number - ) - - iteration.output_data = { - "commit_hash": commit_hash, - "commit_message": commit_message, - "branch": execution.git_manager.feature_branch, - } - - logger.info(f"Changes committed for {execution.task_id}: {commit_hash}") - - async def _step_human_interaction( - self, execution: ExecutionContext, iteration: TaskIteration - ): - """Handle human interaction step""" - # This step waits for human response - the actual waiting is handled - # by the status being WAITING_FOR_HUMAN - iteration.output_data = { - "human_interaction": "waiting_for_response", - "question": iteration.human_question, - } - - async def _step_finalize_task( - self, execution: ExecutionContext, iteration: TaskIteration - ): - """Finalize the task execution""" - execution.status = TaskStatus.COMPLETED - execution.completed_at = datetime.now() - - # Create pull request if Git is configured - pr_url = None - if execution.git_manager: - try: - pr = await execution.git_manager.create_pull_request( - title=f"🤖 {execution.task_data.get('title', 'Task completion')}", - description=f"Autonomous completion of task: {execution.task_data.get('description', '')}", - ) - pr_url = pr.url if pr else None - except Exception as e: - logger.error(f"Failed to create PR for {execution.task_id}: {e}") - - # Prepare final result - execution.result = { - "status": "completed", - "iterations": len(execution.iterations), - "started_at": execution.started_at.isoformat(), - "completed_at": execution.completed_at.isoformat(), - "pull_request_url": pr_url, - "sandbox_id": execution.sandbox.sandbox_id if execution.sandbox else None, - "git_branch": ( - execution.git_manager.feature_branch if execution.git_manager else None - ), - } - - # Update database - await self._update_task_status( - execution.task_id, TaskStatus.COMPLETED, result=execution.result - ) - - iteration.output_data = execution.result - - # End conversation session - if execution.claude_wrapper: - try: - await execution.claude_wrapper.end_conversation_session() - except Exception as e: - logger.error(f"Error ending conversation session: {e}") - - # Store final performance metrics - await self._store_completion_metrics(execution) - - # Extract knowledge from completed task - if self.knowledge_extractor: - try: - extracted_knowledge_ids = ( - await self.knowledge_extractor.extract_knowledge_from_task( - task_id=execution.task_id, - agent_id=execution.agent_id, - execution_result=execution.result, - ) - ) - if extracted_knowledge_ids: - logger.info( - f"Extracted {len(extracted_knowledge_ids)} knowledge items from task {execution.task_id}" - ) - except Exception as e: - logger.error( - f"Error extracting knowledge from task {execution.task_id}: {e}" - ) - - logger.info(f"✅ Task execution completed: {execution.task_id}") - - # End Claude SDK session - if execution.claude_sdk_manager and execution.claude_session_id: - try: - await execution.claude_sdk_manager.terminate_session( - execution.claude_session_id - ) - except Exception as e: - logger.error(f"Error terminating Claude SDK session: {e}") - - # Schedule cleanup - asyncio.create_task(self._cleanup_execution(execution.task_id)) - - def _is_task_complete(self, execution: ExecutionContext) -> bool: - """Check if the task is complete""" - # Simple completion check - in full implementation this would be more sophisticated - return execution.current_iteration >= 3 - - async def _handle_execution_error(self, task_id: str, error_message: str): - """Handle execution error""" - execution = self.active_executions.get(task_id) - if execution: - execution.status = TaskStatus.FAILED - execution.completed_at = datetime.now() - execution.error = error_message - - await self._update_task_status(task_id, TaskStatus.FAILED, error=error_message) - - # Schedule cleanup - asyncio.create_task(self._cleanup_execution(task_id)) - - logger.error(f"Task execution failed: {task_id} - {error_message}") - - async def _cleanup_execution(self, task_id: str): - """Clean up execution resources""" - execution = self.active_executions.pop(task_id, None) - if not execution: - return - - # Cleanup Claude SDK session - if execution.claude_sdk_manager and execution.claude_session_id: - try: - await execution.claude_sdk_manager.terminate_session( - execution.claude_session_id - ) - except Exception as e: - logger.error(f"Error terminating Claude SDK session for {task_id}: {e}") - - # Remove engines from tracking - self.file_operations_engines.pop(task_id, None) - self.claude_sdk_managers.pop(task_id, None) - - # Cleanup sandbox - if execution.sandbox: - try: - await self.sandbox_manager.destroy_sandbox(execution.sandbox.sandbox_id) - except Exception as e: - logger.error(f"Error destroying sandbox for {task_id}: {e}") - - # Cleanup Git workspace (optional - might want to keep for review) - if execution.git_manager and execution.status != TaskStatus.COMPLETED: - try: - await execution.git_manager.cleanup_workspace() - except Exception as e: - logger.error(f"Error cleaning up Git workspace for {task_id}: {e}") - - logger.info(f"Execution cleanup complete: {task_id}") - - async def _monitoring_worker(self): - """Monitor execution health and timeouts""" - while self.running: - try: - current_time = datetime.now() - - for task_id, execution in list(self.active_executions.items()): - # Check for timeouts - if execution.status == TaskStatus.WAITING_FOR_HUMAN: - # Check human response timeout - last_iteration = ( - execution.iterations[-1] if execution.iterations else None - ) - if last_iteration and last_iteration.started_at: - time_waiting = current_time - last_iteration.started_at - if time_waiting > timedelta( - seconds=self.human_response_timeout - ): - await self._handle_execution_error( - task_id, - "Human response timeout - no response received within 24 hours", - ) - else: - # Check general execution timeout - execution_time = current_time - execution.started_at - if execution_time > timedelta( - seconds=self.iteration_timeout * self.max_iterations - ): - await self._handle_execution_error( - task_id, - f"Execution timeout - exceeded maximum time limit", - ) - - # Check iteration limits - if execution.current_iteration > self.max_iterations: - await self._handle_execution_error( - task_id, - f"Maximum iterations exceeded ({self.max_iterations})", - ) - - await asyncio.sleep(60) # Check every minute - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Error in monitoring worker: {e}") - await asyncio.sleep(60) - - async def _cleanup_worker(self): - """Clean up completed executions periodically""" - while self.running: - try: - current_time = datetime.now() - cutoff_time = current_time - timedelta( - hours=1 - ) # Keep completed tasks for 1 hour - - completed_tasks = [ - task_id - for task_id, execution in self.active_executions.items() - if execution.status - in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED] - and execution.completed_at - and execution.completed_at < cutoff_time - ] - - for task_id in completed_tasks: - await self._cleanup_execution(task_id) - - await asyncio.sleep(3600) # Run every hour - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Error in cleanup worker: {e}") - await asyncio.sleep(3600) - - # Database operations - - async def _get_task_data(self, task_id: str) -> Optional[Dict[str, Any]]: - """Get task data from database""" - async with get_db_connection() as conn: - row = await conn.fetchrow("SELECT * FROM tasks WHERE id = $1", task_id) - return dict(row) if row else None - - async def _get_agent_data(self, agent_id: str) -> Optional[Dict[str, Any]]: - """Get agent data from database""" - return await DatabaseManager.get_agent(agent_id) - - async def _update_task_status( - self, - task_id: str, - status: TaskStatus, - result: Optional[Dict[str, Any]] = None, - error: Optional[str] = None, - ): - """Update task status in database""" - await DatabaseManager.update_task_status( - task_id=task_id, status=status.value, result=result - ) - - async def _store_task_iteration(self, task_id: str, iteration: TaskIteration): - """Store task iteration in database""" - async with get_db_connection() as conn: - await conn.execute( - """ - INSERT INTO task_iterations ( - id, task_id, iteration_number, step, started_at, completed_at, - input_data, output_data, success, error_message, - human_question, human_response - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) - """, - str(uuid.uuid4()), - task_id, - iteration.iteration_number, - iteration.step.value, - iteration.started_at, - iteration.completed_at, - json.dumps(iteration.input_data), - json.dumps(iteration.output_data) if iteration.output_data else None, - iteration.success, - iteration.error_message, - iteration.human_question, - iteration.human_response, - ) - - async def _get_task_iterations_from_db(self, task_id: str) -> List[TaskIteration]: - """Get task iterations from database""" - async with get_db_connection() as conn: - rows = await conn.fetch( - """ - SELECT * FROM task_iterations - WHERE task_id = $1 - ORDER BY iteration_number - """, - task_id, - ) - - iterations = [] - for row in rows: - iterations.append( - TaskIteration( - iteration_number=row["iteration_number"], - step=ExecutionStep(row["step"]), - started_at=row["started_at"], - completed_at=row["completed_at"], - input_data=( - json.loads(row["input_data"]) if row["input_data"] else {} - ), - output_data=( - json.loads(row["output_data"]) - if row["output_data"] - else None - ), - success=row["success"], - error_message=row["error_message"], - human_question=row["human_question"], - human_response=row["human_response"], - ) - ) - - return iterations - - async def _store_completion_metrics(self, execution: ExecutionContext): - """Store performance metrics for completed task""" - try: - if not execution.completed_at or not execution.started_at: - return - - # Calculate execution time - execution_time = ( - execution.completed_at - execution.started_at - ).total_seconds() / 60 # minutes - - # Store metrics - await self.conversation_manager.store_performance_metric( - agent_id=execution.agent_id, - task_id=execution.task_id, - metric_type="execution_time_minutes", - metric_value=execution_time, - metric_unit="minutes", - ) - - await self.conversation_manager.store_performance_metric( - agent_id=execution.agent_id, - task_id=execution.task_id, - metric_type="iterations_to_completion", - metric_value=float(execution.current_iteration), - metric_unit="iterations", - ) - - # Store success/failure metric - success_value = 1.0 if execution.status == TaskStatus.COMPLETED else 0.0 - await self.conversation_manager.store_performance_metric( - agent_id=execution.agent_id, - task_id=execution.task_id, - metric_type="task_success_rate", - metric_value=success_value, - metric_unit="boolean", - ) - - except Exception as e: - logger.error(f"Error storing completion metrics: {e}") - - async def _handle_claude_interaction( - self, execution: ExecutionContext, session: ClaudeSDKSession, interaction - ): - """Handle interaction from Claude SDK""" - logger.info( - f"Claude interaction for task {execution.task_id}: {interaction.interaction_type}" - ) - - # This will be processed in the next iteration of _step_execute_iteration - # The interaction handling is done there to maintain the execution flow - - def _format_diffs_for_human(self, diffs: Dict[str, str]) -> str: - """Format file diffs for human review""" - if not diffs: - return "No file changes detected." - - formatted = [] - for file_path, diff in diffs.items(): - formatted.append(f"\n--- {file_path} ---") - formatted.append(diff[:1000] + "..." if len(diff) > 1000 else diff) - - return "\n".join(formatted) - - async def handle_human_response(self, task_id: str, response: str) -> bool: - """Enhanced human response handler that integrates with Claude SDK""" - - execution = self.active_executions.get(task_id) - if not execution or execution.status != TaskStatus.WAITING_FOR_HUMAN: - return False - - # Find the current iteration waiting for human response - current_iteration = None - for iteration in reversed(execution.iterations): - if iteration.human_question and not iteration.human_response: - current_iteration = iteration - break - - if current_iteration: - current_iteration.human_response = response - execution.status = TaskStatus.EXECUTING - - # Handle different types of responses - output_data = current_iteration.output_data or {} - - if output_data.get("file_approval_required"): - # Handle file approval - batch_id = output_data.get("batch_id") - approved = response.lower().strip() in [ - "yes", - "y", - "approve", - "approved", - "true", - ] - - if ( - batch_id - and execution.claude_sdk_manager - and execution.claude_session_id - ): - success = ( - await execution.claude_sdk_manager.approve_file_operations( - execution.claude_session_id, batch_id, approved - ) - ) - - if success: - logger.info( - f"File operations {'approved' if approved else 'rejected'} for task {task_id}" - ) - else: - logger.error( - f"Failed to process file approval for task {task_id}" - ) - - else: - # Handle general user input - if execution.claude_sdk_manager and execution.claude_session_id: - success = await execution.claude_sdk_manager.send_input( - execution.claude_session_id, response - ) - - if success: - logger.info( - f"Human response sent to Claude SDK for task {task_id}" - ) - else: - logger.error( - f"Failed to send human response to Claude SDK for task {task_id}" - ) - - # Update database - await self._store_task_iteration(task_id, current_iteration) - await self._update_task_status(task_id, TaskStatus.EXECUTING) - - logger.info(f"Human response processed for task {task_id}") - return True - - return False +""" +Task Execution Engine for FuzeAgent Autonomous Execution + +Orchestrates the autonomous execution of tasks by agents, managing: +- Task lifecycle and state transitions +- Sandbox creation and cleanup +- Git workflow automation +- Human-in-the-loop interactions +- Inter-agent communication +- Result aggregation + +This is the core component that ties together all autonomous execution components. +""" + +import asyncio +import json +import logging +import uuid +from dataclasses import dataclass +from datetime import datetime, timedelta +from enum import Enum +from typing import Any, Callable, Dict, List, Optional + +from .claude_code_wrapper import ClaudeCodeWrapper +from .claude_sdk_manager import ClaudeSDKManager, ClaudeSDKSession +from .context_enhancement_service import ContextEnhancementService +from .conversation_manager import ConversationManager, InteractionType +from .database import DatabaseManager, get_db_connection +from .file_operations_engine import FileOperationsEngine +from .git_workflow_manager import GitWorkflowManager +from .sandbox_manager import AgentSandboxManager, Sandbox +from .task_knowledge_extractor import TaskKnowledgeExtractor + +logger = logging.getLogger(__name__) + + +class TaskStatus(str, Enum): + PENDING = "pending" + ANALYZING = "analyzing" + SETTING_UP = "setting_up" + EXECUTING = "executing" + WAITING_FOR_HUMAN = "waiting_for_human" + REVIEWING = "reviewing" + COMMITTING = "committing" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +class ExecutionStep(str, Enum): + ANALYZE_TASK = "analyze_task" + SETUP_SANDBOX = "setup_sandbox" + SETUP_GIT = "setup_git" + EXECUTE_ITERATION = "execute_iteration" + REVIEW_CHANGES = "review_changes" + COMMIT_CHANGES = "commit_changes" + HUMAN_INTERACTION = "human_interaction" + FINALIZE_TASK = "finalize_task" + CLEANUP = "cleanup" + + +@dataclass +class TaskIteration: + """Represents a single iteration of task execution""" + + iteration_number: int + step: ExecutionStep + started_at: datetime + completed_at: Optional[datetime] + input_data: Dict[str, Any] + output_data: Optional[Dict[str, Any]] + success: bool + error_message: Optional[str] + human_question: Optional[str] = None + human_response: Optional[str] = None + + +@dataclass +class ExecutionContext: + """Context for task execution""" + + task_id: str + agent_id: str + task_data: Dict[str, Any] + agent_data: Dict[str, Any] + sandbox: Optional[Sandbox] + git_manager: Optional[GitWorkflowManager] + claude_wrapper: Optional[ClaudeCodeWrapper] + current_iteration: int + iterations: List[TaskIteration] + status: TaskStatus + started_at: datetime + completed_at: Optional[datetime] + result: Optional[Dict[str, Any]] + error: Optional[str] + # New components for autonomous execution + file_operations_engine: Optional[FileOperationsEngine] = None + claude_sdk_manager: Optional[ClaudeSDKManager] = None + claude_session_id: Optional[str] = None + + +class TaskExecutionEngine: + """ + Orchestrates autonomous task execution by agents. + + Features: + - Task lifecycle management + - Sandbox and Git workflow integration + - Human-in-the-loop interactions + - Dependency handling + - Result aggregation + - Error recovery + """ + + def __init__( + self, + sandbox_manager: AgentSandboxManager, + knowledge_extractor: Optional[TaskKnowledgeExtractor] = None, + context_enhancer: Optional[ContextEnhancementService] = None, + ): + self.sandbox_manager = sandbox_manager + self.conversation_manager = ConversationManager() + self.active_executions: Dict[str, ExecutionContext] = {} + self.execution_callbacks: Dict[str, List[Callable]] = {} + self.running = False + self.worker_tasks: List[asyncio.Task] = [] + + # Knowledge management services + self.knowledge_extractor = knowledge_extractor + self.context_enhancer = context_enhancer + + # Initialize integrated components + self.file_operations_engines: Dict[str, FileOperationsEngine] = {} # Per task + self.claude_sdk_managers: Dict[str, ClaudeSDKManager] = {} # Per task + + # Configuration + self.max_iterations = 50 + self.iteration_timeout = 3600 # 1 hour per iteration + self.human_response_timeout = 86400 # 24 hours for human response + + async def start(self): + """Start the execution engine""" + logger.info("Starting TaskExecutionEngine") + self.running = True + + # Start worker tasks + self.worker_tasks = [ + asyncio.create_task(self._execution_worker()), + asyncio.create_task(self._monitoring_worker()), + asyncio.create_task(self._cleanup_worker()), + ] + + logger.info("TaskExecutionEngine started") + + async def stop(self): + """Stop the execution engine""" + logger.info("Stopping TaskExecutionEngine") + self.running = False + + # Cancel worker tasks + for task in self.worker_tasks: + task.cancel() + + try: + await asyncio.gather(*self.worker_tasks, return_exceptions=True) + except Exception as e: + logger.error(f"Error stopping worker tasks: {e}") + + # Clean up active executions + for execution_id in list(self.active_executions.keys()): + try: + await self._cleanup_execution(execution_id) + except Exception as e: + logger.error(f"Error cleaning up execution {execution_id}: {e}") + + logger.info("TaskExecutionEngine stopped") + + async def start_task_execution(self, task_id: str) -> Dict[str, Any]: + """ + Start autonomous execution of a task. + Returns execution status and context. + """ + logger.info(f"Starting task execution: {task_id}") + + try: + # Get task data + task_data = await self._get_task_data(task_id) + if not task_data: + raise ValueError(f"Task {task_id} not found") + + # Get agent data + agent_id = task_data.get("assigned_to") + if not agent_id: + raise ValueError(f"Task {task_id} has no assigned agent") + + agent_data = await self._get_agent_data(agent_id) + if not agent_data: + raise ValueError(f"Agent {agent_id} not found") + + # Create execution context + execution_context = ExecutionContext( + task_id=task_id, + agent_id=agent_id, + task_data=task_data, + agent_data=agent_data, + sandbox=None, + git_manager=None, + claude_wrapper=None, + current_iteration=0, + iterations=[], + status=TaskStatus.PENDING, + started_at=datetime.now(), + completed_at=None, + result=None, + error=None, + ) + + # Store execution context + self.active_executions[task_id] = execution_context + + # Update task status in database + await self._update_task_status(task_id, TaskStatus.PENDING) + + logger.info(f"✅ Task execution started: {task_id}") + return { + "task_id": task_id, + "status": TaskStatus.PENDING.value, + "execution_started": True, + "agent_id": agent_id, + } + + except Exception as e: + logger.error(f"❌ Failed to start task execution {task_id}: {e}") + await self._update_task_status(task_id, TaskStatus.FAILED, error=str(e)) + raise + + async def get_execution_status(self, task_id: str) -> Dict[str, Any]: + """Get detailed execution status for a task""" + + execution = self.active_executions.get(task_id) + if not execution: + # Check database for completed/failed tasks + task_data = await self._get_task_data(task_id) + if task_data: + return { + "task_id": task_id, + "status": task_data.get("status", "unknown"), + "result": task_data.get("result"), + "active_execution": False, + } + else: + return {"task_id": task_id, "status": "not_found"} + + return { + "task_id": task_id, + "status": execution.status.value, + "agent_id": execution.agent_id, + "current_iteration": execution.current_iteration, + "iterations_count": len(execution.iterations), + "started_at": execution.started_at.isoformat(), + "completed_at": ( + execution.completed_at.isoformat() if execution.completed_at else None + ), + "sandbox_id": execution.sandbox.sandbox_id if execution.sandbox else None, + "git_branch": ( + execution.git_manager.feature_branch if execution.git_manager else None + ), + "result": execution.result, + "error": execution.error, + "active_execution": True, + } + + async def get_task_iterations(self, task_id: str) -> List[Dict[str, Any]]: + """Get iteration history for a task""" + + execution = self.active_executions.get(task_id) + if execution: + iterations = execution.iterations + else: + # Get from database + iterations = await self._get_task_iterations_from_db(task_id) + + return [ + { + "iteration_number": it.iteration_number, + "step": it.step.value if hasattr(it.step, "value") else str(it.step), + "started_at": it.started_at.isoformat(), + "completed_at": ( + it.completed_at.isoformat() if it.completed_at else None + ), + "success": it.success, + "error_message": it.error_message, + "human_question": it.human_question, + "human_response": it.human_response, + "input_data": it.input_data, + "output_data": it.output_data, + } + for it in iterations + ] + + async def cancel_task_execution(self, task_id: str) -> bool: + """Cancel a running task execution""" + + execution = self.active_executions.get(task_id) + if not execution: + return False + + execution.status = TaskStatus.CANCELLED + execution.completed_at = datetime.now() + execution.error = "Task cancelled by user" + + # Update database + await self._update_task_status( + task_id, TaskStatus.CANCELLED, error="Task cancelled by user" + ) + + # Schedule cleanup + asyncio.create_task(self._cleanup_execution(task_id)) + + logger.info(f"Task execution cancelled: {task_id}") + return True + + # Private methods for execution workflow + + async def _execution_worker(self): + """Main execution worker that processes pending tasks""" + while self.running: + try: + # Find tasks ready for execution + pending_tasks = [ + task_id + for task_id, execution in self.active_executions.items() + if execution.status in [TaskStatus.PENDING, TaskStatus.EXECUTING] + ] + + # Process each pending task + for task_id in pending_tasks: + try: + await self._process_task_execution(task_id) + except Exception as e: + logger.error(f"Error processing task {task_id}: {e}") + await self._handle_execution_error(task_id, str(e)) + + # Sleep between iterations + await asyncio.sleep(5) + + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"Error in execution worker: {e}") + await asyncio.sleep(10) + + async def _process_task_execution(self, task_id: str): + """Process a single task execution step""" + execution = self.active_executions.get(task_id) + if not execution: + return + + # Skip if waiting for human or in terminal state + if execution.status in [ + TaskStatus.WAITING_FOR_HUMAN, + TaskStatus.COMPLETED, + TaskStatus.FAILED, + TaskStatus.CANCELLED, + ]: + return + + # Determine next step + next_step = self._determine_next_step(execution) + if not next_step: + return + + # Execute the step + try: + await self._execute_step(execution, next_step) + except Exception as e: + logger.error(f"Error executing step {next_step} for task {task_id}: {e}") + await self._handle_execution_error(task_id, str(e)) + + def _determine_next_step( + self, execution: ExecutionContext + ) -> Optional[ExecutionStep]: + """Determine the next execution step""" + + if execution.status == TaskStatus.PENDING: + return ExecutionStep.ANALYZE_TASK + + if not execution.iterations: + return ExecutionStep.ANALYZE_TASK + + last_iteration = execution.iterations[-1] + + # Continue based on last completed step + if last_iteration.step == ExecutionStep.ANALYZE_TASK and last_iteration.success: + return ExecutionStep.SETUP_SANDBOX + elif ( + last_iteration.step == ExecutionStep.SETUP_SANDBOX + and last_iteration.success + ): + return ExecutionStep.SETUP_GIT + elif last_iteration.step == ExecutionStep.SETUP_GIT and last_iteration.success: + return ExecutionStep.EXECUTE_ITERATION + elif ( + last_iteration.step == ExecutionStep.EXECUTE_ITERATION + and last_iteration.success + ): + # Check if we need human input + if last_iteration.human_question: + return ExecutionStep.HUMAN_INTERACTION + else: + return ExecutionStep.REVIEW_CHANGES + elif ( + last_iteration.step == ExecutionStep.HUMAN_INTERACTION + and last_iteration.human_response + ): + return ExecutionStep.EXECUTE_ITERATION + elif ( + last_iteration.step == ExecutionStep.REVIEW_CHANGES + and last_iteration.success + ): + return ExecutionStep.COMMIT_CHANGES + elif ( + last_iteration.step == ExecutionStep.COMMIT_CHANGES + and last_iteration.success + ): + # Check if task is complete + if self._is_task_complete(execution): + return ExecutionStep.FINALIZE_TASK + else: + return ExecutionStep.EXECUTE_ITERATION + + return None + + async def _execute_step(self, execution: ExecutionContext, step: ExecutionStep): + """Execute a specific step""" + + iteration = TaskIteration( + iteration_number=execution.current_iteration + 1, + step=step, + started_at=datetime.now(), + completed_at=None, + input_data={}, + output_data=None, + success=False, + error_message=None, + ) + + execution.iterations.append(iteration) + execution.current_iteration += 1 + + try: + if step == ExecutionStep.ANALYZE_TASK: + await self._step_analyze_task(execution, iteration) + elif step == ExecutionStep.SETUP_SANDBOX: + await self._step_setup_sandbox(execution, iteration) + elif step == ExecutionStep.SETUP_GIT: + await self._step_setup_git(execution, iteration) + elif step == ExecutionStep.EXECUTE_ITERATION: + await self._step_execute_iteration(execution, iteration) + elif step == ExecutionStep.REVIEW_CHANGES: + await self._step_review_changes(execution, iteration) + elif step == ExecutionStep.COMMIT_CHANGES: + await self._step_commit_changes(execution, iteration) + elif step == ExecutionStep.HUMAN_INTERACTION: + await self._step_human_interaction(execution, iteration) + elif step == ExecutionStep.FINALIZE_TASK: + await self._step_finalize_task(execution, iteration) + + iteration.completed_at = datetime.now() + iteration.success = True + + except Exception as e: + iteration.completed_at = datetime.now() + iteration.success = False + iteration.error_message = str(e) + raise + + finally: + # Store iteration in database + await self._store_task_iteration(execution.task_id, iteration) + + async def _step_analyze_task( + self, execution: ExecutionContext, iteration: TaskIteration + ): + """Analyze the task and prepare execution plan""" + execution.status = TaskStatus.ANALYZING + await self._update_task_status(execution.task_id, TaskStatus.ANALYZING) + + # Analyze task requirements + task_description = execution.task_data.get("description", "") + task_title = execution.task_data.get("title", "") + + iteration.input_data = { + "task_title": task_title, + "task_description": task_description, + "agent_type": execution.agent_data.get("type"), + "agent_role": execution.agent_data.get("role"), + } + + # Enhance context with organizational knowledge + enhanced_context = None + if self.context_enhancer: + try: + enhanced_context = await self.context_enhancer.enhance_agent_context( + agent_id=execution.agent_id, + task_data=execution.task_data, + base_context=iteration.input_data, + ) + logger.info( + f"Enhanced context for task {execution.task_id}: " + f"{len(enhanced_context.organizational_knowledge)} org + " + f"{len(enhanced_context.team_knowledge)} team + " + f"{len(enhanced_context.similar_task_insights)} similar task insights" + ) + except Exception as e: + logger.error( + f"Error enhancing context for task {execution.task_id}: {e}" + ) + + # Simple analysis for now - in a full implementation this would use AI + iteration.output_data = { + "analysis_complete": True, + "requires_sandbox": execution.agent_data.get("type") == "developer", + "requires_git": bool( + execution.agent_data.get("repository_settings", {}).get( + "repository_url" + ) + ), + "enhanced_context": enhanced_context, + "estimated_complexity": "medium", + "estimated_iterations": 5, + } + + logger.info(f"Task analysis complete for {execution.task_id}") + + async def _step_setup_sandbox( + self, execution: ExecutionContext, iteration: TaskIteration + ): + """Set up sandbox environment for the agent""" + execution.status = TaskStatus.SETTING_UP + await self._update_task_status(execution.task_id, TaskStatus.SETTING_UP) + + agent_template = execution.agent_data.get("template_id", "python_developer") + repository_settings = execution.agent_data.get("repository_settings", {}) + sandbox_settings = execution.agent_data.get("sandbox_settings", {}) + + # Create sandbox + sandbox = await self.sandbox_manager.create_sandbox( + agent_id=execution.agent_id, + task_id=execution.task_id, + agent_template=agent_template, + repository_settings=repository_settings, + custom_settings=sandbox_settings, + ) + + execution.sandbox = sandbox + + iteration.input_data = { + "agent_template": agent_template, + "repository_settings": repository_settings, + "sandbox_settings": sandbox_settings, + } + + iteration.output_data = { + "sandbox_id": sandbox.sandbox_id, + "workspace_path": sandbox.workspace_path, + "container_id": sandbox.container_id, + } + + logger.info( + f"Sandbox setup complete for {execution.task_id}: {sandbox.sandbox_id}" + ) + + async def _step_setup_git( + self, execution: ExecutionContext, iteration: TaskIteration + ): + """Set up Git workflow for the task""" + repository_settings = execution.agent_data.get("repository_settings", {}) + + if not repository_settings.get("repository_url"): + # Skip Git setup if no repository + iteration.output_data = { + "git_setup": "skipped", + "reason": "no_repository_configured", + } + return + + # Create Git workflow manager + git_manager = GitWorkflowManager( + agent_id=execution.agent_id, + task_id=execution.task_id, + repo_settings=repository_settings, + ) + + # Setup workspace + feature_branch = await git_manager.setup_workspace() + + execution.git_manager = git_manager + + # Create enhanced Claude wrapper with Git context and conversation tracking + execution.claude_wrapper = ClaudeCodeWrapper( + workspace_path=git_manager.workspace_path, + git_manager=git_manager, + agent_id=execution.agent_id, + task_id=execution.task_id, + conversation_manager=self.conversation_manager, + ) + + # Initialize File Operations Engine + file_ops_engine = FileOperationsEngine(git_manager.workspace_path) + execution.file_operations_engine = file_ops_engine + self.file_operations_engines[execution.task_id] = file_ops_engine + + # Initialize Claude SDK Manager + claude_sdk_manager = ClaudeSDKManager( + file_operations_engine=file_ops_engine, + conversation_manager=self.conversation_manager, + ) + execution.claude_sdk_manager = claude_sdk_manager + self.claude_sdk_managers[execution.task_id] = claude_sdk_manager + + # Start conversation session + await execution.claude_wrapper.start_conversation_session( + execution.sandbox.sandbox_id + ) + + iteration.input_data = { + "repository_url": repository_settings.get("repository_url"), + "default_branch": repository_settings.get("default_branch", "main"), + } + + iteration.output_data = { + "git_setup": "complete", + "feature_branch": feature_branch, + "workspace_path": git_manager.workspace_path, + } + + logger.info(f"Git setup complete for {execution.task_id}: {feature_branch}") + + async def _step_execute_iteration( + self, execution: ExecutionContext, iteration: TaskIteration + ): + """Execute a development iteration using Claude SDK Manager""" + execution.status = TaskStatus.EXECUTING + await self._update_task_status(execution.task_id, TaskStatus.EXECUTING) + + task_description = execution.task_data.get("description", "") + task_title = execution.task_data.get("title", "") + + iteration.input_data = { + "task_description": task_description, + "task_title": task_title, + "iteration_number": iteration.iteration_number, + "workspace_path": ( + execution.sandbox.workspace_path if execution.sandbox else None + ), + } + + try: + # Start Claude SDK session if not already running + if not execution.claude_session_id and execution.claude_sdk_manager: + # Build comprehensive task context + context_info = "" + if execution.git_manager: + context_info += f"\nRepository: {execution.git_manager.repo_url}" + context_info += f"\nBranch: {execution.git_manager.feature_branch}" + if iteration.iteration_number > 1: + context_info += f"\nIteration: {iteration.iteration_number} of ongoing development" + + # Construct task prompt + task_prompt = f""" +Task: {task_title} + +Description: {task_description} + +Context: {context_info} + +Please analyze the codebase, understand the requirements, and implement the necessary changes. +Work incrementally and ask for clarification if needed. +""" + + # Start Claude SDK session + execution.claude_session_id = ( + await execution.claude_sdk_manager.start_session( + task_id=execution.task_id, + agent_id=execution.agent_id, + workspace_path=( + execution.git_manager.workspace_path + if execution.git_manager + else execution.sandbox.workspace_path + ), + task_description=task_prompt, + additional_context=context_info, + ) + ) + + logger.info( + f"Started Claude SDK session: {execution.claude_session_id}" + ) + + # Register interaction callback for human-in-the-loop + if execution.claude_sdk_manager and execution.claude_session_id: + execution.claude_sdk_manager.register_interaction_callback( + execution.claude_session_id, + lambda session, interaction: self._handle_claude_interaction( + execution, session, interaction + ), + ) + + # Monitor session status + session_status = ( + await execution.claude_sdk_manager.get_session_status( + execution.claude_session_id + ) + if execution.claude_session_id + else None + ) + + if session_status: + current_interaction = session_status.get("current_interaction") + if current_interaction: + # Claude is waiting for human input + interaction_type = current_interaction.get("type") + if interaction_type in ["user_input", "confirmation"]: + iteration.human_question = current_interaction.get("prompt") + execution.status = TaskStatus.WAITING_FOR_HUMAN + await self._update_task_status( + execution.task_id, TaskStatus.WAITING_FOR_HUMAN + ) + + iteration.output_data = { + "claude_session_status": session_status.get("state"), + "human_interaction_required": True, + "interaction_type": interaction_type, + "human_question_asked": True, + } + + logger.info( + f"Claude SDK requesting human input for task {execution.task_id}" + ) + return + + elif interaction_type == "file_approval": + # File operations pending approval + batch_id = current_interaction.get("metadata", {}).get( + "batch_id" + ) + if batch_id and execution.file_operations_engine: + # Get diff preview for human review + diffs = await execution.file_operations_engine.get_file_diff_preview( + batch_id + ) + + iteration.human_question = f"""Claude wants to make the following file changes: + +{current_interaction.get('prompt')} + +File changes preview: +{self._format_diffs_for_human(diffs)} + +Approve these changes? (yes/no)""" + + execution.status = TaskStatus.WAITING_FOR_HUMAN + await self._update_task_status( + execution.task_id, TaskStatus.WAITING_FOR_HUMAN + ) + + iteration.output_data = { + "claude_session_status": session_status.get("state"), + "file_approval_required": True, + "batch_id": batch_id, + "file_changes_preview": diffs, + "human_question_asked": True, + } + + logger.info( + f"Claude SDK requesting file approval for task {execution.task_id}" + ) + return + + # No interaction needed - continue execution + iteration.output_data = { + "claude_session_status": session_status.get("state"), + "session_active": True, + "iteration_completed": True, + "workspace_path": session_status.get("workspace_path"), + } + + # Check if session completed + if session_status.get("state") in ["completed", "terminated"]: + iteration.output_data["development_complete"] = True + + else: + # No session - this shouldn't happen but handle gracefully + iteration.output_data = { + "error": "No Claude SDK session available", + "development_complete": False, + } + + except Exception as e: + logger.error(f"Error in Claude SDK iteration: {e}") + iteration.output_data = {"error": str(e), "development_complete": False} + raise + + logger.info( + f"Development iteration {iteration.iteration_number} processed for {execution.task_id}" + ) + + async def _step_review_changes( + self, execution: ExecutionContext, iteration: TaskIteration + ): + """Review the changes made in the iteration""" + execution.status = TaskStatus.REVIEWING + await self._update_task_status(execution.task_id, TaskStatus.REVIEWING) + + # Review changes - in full implementation this would: + # 1. Run linting and type checking + # 2. Run tests + # 3. Check code quality + # 4. Validate against requirements + + iteration.output_data = { + "review_passed": True, + "issues_found": [], + "tests_passed": True, + "code_quality_score": 85, + } + + logger.info(f"Code review complete for {execution.task_id}") + + async def _step_commit_changes( + self, execution: ExecutionContext, iteration: TaskIteration + ): + """Commit changes to Git""" + execution.status = TaskStatus.COMMITTING + await self._update_task_status(execution.task_id, TaskStatus.COMMITTING) + + if not execution.git_manager: + iteration.output_data = {"commit": "skipped", "reason": "no_git_manager"} + return + + # Commit changes + commit_message = f"Iteration {iteration.iteration_number}: {execution.task_data.get('title', 'Task update')}" + commit_hash = await execution.git_manager.commit_changes( + message=commit_message, iteration_number=iteration.iteration_number + ) + + iteration.output_data = { + "commit_hash": commit_hash, + "commit_message": commit_message, + "branch": execution.git_manager.feature_branch, + } + + logger.info(f"Changes committed for {execution.task_id}: {commit_hash}") + + async def _step_human_interaction( + self, execution: ExecutionContext, iteration: TaskIteration + ): + """Handle human interaction step""" + # This step waits for human response - the actual waiting is handled + # by the status being WAITING_FOR_HUMAN + iteration.output_data = { + "human_interaction": "waiting_for_response", + "question": iteration.human_question, + } + + async def _step_finalize_task( + self, execution: ExecutionContext, iteration: TaskIteration + ): + """Finalize the task execution""" + execution.status = TaskStatus.COMPLETED + execution.completed_at = datetime.now() + + # Create pull request if Git is configured + pr_url = None + if execution.git_manager: + try: + pr = await execution.git_manager.create_pull_request( + title=f"🤖 {execution.task_data.get('title', 'Task completion')}", + description=f"Autonomous completion of task: {execution.task_data.get('description', '')}", + ) + pr_url = pr.url if pr else None + except Exception as e: + logger.error(f"Failed to create PR for {execution.task_id}: {e}") + + # Prepare final result + execution.result = { + "status": "completed", + "iterations": len(execution.iterations), + "started_at": execution.started_at.isoformat(), + "completed_at": execution.completed_at.isoformat(), + "pull_request_url": pr_url, + "sandbox_id": execution.sandbox.sandbox_id if execution.sandbox else None, + "git_branch": ( + execution.git_manager.feature_branch if execution.git_manager else None + ), + } + + # Update database + await self._update_task_status( + execution.task_id, TaskStatus.COMPLETED, result=execution.result + ) + + iteration.output_data = execution.result + + # End conversation session + if execution.claude_wrapper: + try: + await execution.claude_wrapper.end_conversation_session() + except Exception as e: + logger.error(f"Error ending conversation session: {e}") + + # Store final performance metrics + await self._store_completion_metrics(execution) + + # Extract knowledge from completed task + if self.knowledge_extractor: + try: + extracted_knowledge_ids = ( + await self.knowledge_extractor.extract_knowledge_from_task( + task_id=execution.task_id, + agent_id=execution.agent_id, + execution_result=execution.result, + ) + ) + if extracted_knowledge_ids: + logger.info( + f"Extracted {len(extracted_knowledge_ids)} knowledge items from task {execution.task_id}" + ) + except Exception as e: + logger.error( + f"Error extracting knowledge from task {execution.task_id}: {e}" + ) + + logger.info(f"✅ Task execution completed: {execution.task_id}") + + # End Claude SDK session + if execution.claude_sdk_manager and execution.claude_session_id: + try: + await execution.claude_sdk_manager.terminate_session( + execution.claude_session_id + ) + except Exception as e: + logger.error(f"Error terminating Claude SDK session: {e}") + + # Schedule cleanup + asyncio.create_task(self._cleanup_execution(execution.task_id)) + + def _is_task_complete(self, execution: ExecutionContext) -> bool: + """Check if the task is complete""" + # Simple completion check - in full implementation this would be more sophisticated + return execution.current_iteration >= 3 + + async def _handle_execution_error(self, task_id: str, error_message: str): + """Handle execution error""" + execution = self.active_executions.get(task_id) + if execution: + execution.status = TaskStatus.FAILED + execution.completed_at = datetime.now() + execution.error = error_message + + await self._update_task_status(task_id, TaskStatus.FAILED, error=error_message) + + # Schedule cleanup + asyncio.create_task(self._cleanup_execution(task_id)) + + logger.error(f"Task execution failed: {task_id} - {error_message}") + + async def _cleanup_execution(self, task_id: str): + """Clean up execution resources""" + execution = self.active_executions.pop(task_id, None) + if not execution: + return + + # Cleanup Claude SDK session + if execution.claude_sdk_manager and execution.claude_session_id: + try: + await execution.claude_sdk_manager.terminate_session( + execution.claude_session_id + ) + except Exception as e: + logger.error(f"Error terminating Claude SDK session for {task_id}: {e}") + + # Remove engines from tracking + self.file_operations_engines.pop(task_id, None) + self.claude_sdk_managers.pop(task_id, None) + + # Cleanup sandbox + if execution.sandbox: + try: + await self.sandbox_manager.destroy_sandbox(execution.sandbox.sandbox_id) + except Exception as e: + logger.error(f"Error destroying sandbox for {task_id}: {e}") + + # Cleanup Git workspace (optional - might want to keep for review) + if execution.git_manager and execution.status != TaskStatus.COMPLETED: + try: + await execution.git_manager.cleanup_workspace() + except Exception as e: + logger.error(f"Error cleaning up Git workspace for {task_id}: {e}") + + logger.info(f"Execution cleanup complete: {task_id}") + + async def _monitoring_worker(self): + """Monitor execution health and timeouts""" + while self.running: + try: + current_time = datetime.now() + + for task_id, execution in list(self.active_executions.items()): + # Check for timeouts + if execution.status == TaskStatus.WAITING_FOR_HUMAN: + # Check human response timeout + last_iteration = ( + execution.iterations[-1] if execution.iterations else None + ) + if last_iteration and last_iteration.started_at: + time_waiting = current_time - last_iteration.started_at + if time_waiting > timedelta( + seconds=self.human_response_timeout + ): + await self._handle_execution_error( + task_id, + "Human response timeout - no response received within 24 hours", + ) + else: + # Check general execution timeout + execution_time = current_time - execution.started_at + if execution_time > timedelta( + seconds=self.iteration_timeout * self.max_iterations + ): + await self._handle_execution_error( + task_id, + f"Execution timeout - exceeded maximum time limit", + ) + + # Check iteration limits + if execution.current_iteration > self.max_iterations: + await self._handle_execution_error( + task_id, + f"Maximum iterations exceeded ({self.max_iterations})", + ) + + await asyncio.sleep(60) # Check every minute + + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"Error in monitoring worker: {e}") + await asyncio.sleep(60) + + async def _cleanup_worker(self): + """Clean up completed executions periodically""" + while self.running: + try: + current_time = datetime.now() + cutoff_time = current_time - timedelta( + hours=1 + ) # Keep completed tasks for 1 hour + + completed_tasks = [ + task_id + for task_id, execution in self.active_executions.items() + if execution.status + in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED] + and execution.completed_at + and execution.completed_at < cutoff_time + ] + + for task_id in completed_tasks: + await self._cleanup_execution(task_id) + + await asyncio.sleep(3600) # Run every hour + + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"Error in cleanup worker: {e}") + await asyncio.sleep(3600) + + # Database operations + + async def _get_task_data(self, task_id: str) -> Optional[Dict[str, Any]]: + """Get task data from database""" + async with get_db_connection() as conn: + row = await conn.fetchrow("SELECT * FROM tasks WHERE id = $1", task_id) + return dict(row) if row else None + + async def _get_agent_data(self, agent_id: str) -> Optional[Dict[str, Any]]: + """Get agent data from database""" + return await DatabaseManager.get_agent(agent_id) + + async def _update_task_status( + self, + task_id: str, + status: TaskStatus, + result: Optional[Dict[str, Any]] = None, + error: Optional[str] = None, + ): + """Update task status in database""" + await DatabaseManager.update_task_status( + task_id=task_id, status=status.value, result=result + ) + + async def _store_task_iteration(self, task_id: str, iteration: TaskIteration): + """Store task iteration in database""" + async with get_db_connection() as conn: + await conn.execute( + """ + INSERT INTO task_iterations ( + id, task_id, iteration_number, step, started_at, completed_at, + input_data, output_data, success, error_message, + human_question, human_response + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) + """, + str(uuid.uuid4()), + task_id, + iteration.iteration_number, + iteration.step.value, + iteration.started_at, + iteration.completed_at, + json.dumps(iteration.input_data), + json.dumps(iteration.output_data) if iteration.output_data else None, + iteration.success, + iteration.error_message, + iteration.human_question, + iteration.human_response, + ) + + async def _get_task_iterations_from_db(self, task_id: str) -> List[TaskIteration]: + """Get task iterations from database""" + async with get_db_connection() as conn: + rows = await conn.fetch( + """ + SELECT * FROM task_iterations + WHERE task_id = $1 + ORDER BY iteration_number + """, + task_id, + ) + + iterations = [] + for row in rows: + iterations.append( + TaskIteration( + iteration_number=row["iteration_number"], + step=ExecutionStep(row["step"]), + started_at=row["started_at"], + completed_at=row["completed_at"], + input_data=( + json.loads(row["input_data"]) if row["input_data"] else {} + ), + output_data=( + json.loads(row["output_data"]) + if row["output_data"] + else None + ), + success=row["success"], + error_message=row["error_message"], + human_question=row["human_question"], + human_response=row["human_response"], + ) + ) + + return iterations + + async def _store_completion_metrics(self, execution: ExecutionContext): + """Store performance metrics for completed task""" + try: + if not execution.completed_at or not execution.started_at: + return + + # Calculate execution time + execution_time = ( + execution.completed_at - execution.started_at + ).total_seconds() / 60 # minutes + + # Store metrics + await self.conversation_manager.store_performance_metric( + agent_id=execution.agent_id, + task_id=execution.task_id, + metric_type="execution_time_minutes", + metric_value=execution_time, + metric_unit="minutes", + ) + + await self.conversation_manager.store_performance_metric( + agent_id=execution.agent_id, + task_id=execution.task_id, + metric_type="iterations_to_completion", + metric_value=float(execution.current_iteration), + metric_unit="iterations", + ) + + # Store success/failure metric + success_value = 1.0 if execution.status == TaskStatus.COMPLETED else 0.0 + await self.conversation_manager.store_performance_metric( + agent_id=execution.agent_id, + task_id=execution.task_id, + metric_type="task_success_rate", + metric_value=success_value, + metric_unit="boolean", + ) + + except Exception as e: + logger.error(f"Error storing completion metrics: {e}") + + async def _handle_claude_interaction( + self, execution: ExecutionContext, session: ClaudeSDKSession, interaction + ): + """Handle interaction from Claude SDK""" + logger.info( + f"Claude interaction for task {execution.task_id}: {interaction.interaction_type}" + ) + + # This will be processed in the next iteration of _step_execute_iteration + # The interaction handling is done there to maintain the execution flow + + def _format_diffs_for_human(self, diffs: Dict[str, str]) -> str: + """Format file diffs for human review""" + if not diffs: + return "No file changes detected." + + formatted = [] + for file_path, diff in diffs.items(): + formatted.append(f"\n--- {file_path} ---") + formatted.append(diff[:1000] + "..." if len(diff) > 1000 else diff) + + return "\n".join(formatted) + + async def handle_human_response(self, task_id: str, response: str) -> bool: + """Enhanced human response handler that integrates with Claude SDK""" + + execution = self.active_executions.get(task_id) + if not execution or execution.status != TaskStatus.WAITING_FOR_HUMAN: + return False + + # Find the current iteration waiting for human response + current_iteration = None + for iteration in reversed(execution.iterations): + if iteration.human_question and not iteration.human_response: + current_iteration = iteration + break + + if current_iteration: + current_iteration.human_response = response + execution.status = TaskStatus.EXECUTING + + # Handle different types of responses + output_data = current_iteration.output_data or {} + + if output_data.get("file_approval_required"): + # Handle file approval + batch_id = output_data.get("batch_id") + approved = response.lower().strip() in [ + "yes", + "y", + "approve", + "approved", + "true", + ] + + if ( + batch_id + and execution.claude_sdk_manager + and execution.claude_session_id + ): + success = ( + await execution.claude_sdk_manager.approve_file_operations( + execution.claude_session_id, batch_id, approved + ) + ) + + if success: + logger.info( + f"File operations {'approved' if approved else 'rejected'} for task {task_id}" + ) + else: + logger.error( + f"Failed to process file approval for task {task_id}" + ) + + else: + # Handle general user input + if execution.claude_sdk_manager and execution.claude_session_id: + success = await execution.claude_sdk_manager.send_input( + execution.claude_session_id, response + ) + + if success: + logger.info( + f"Human response sent to Claude SDK for task {task_id}" + ) + else: + logger.error( + f"Failed to send human response to Claude SDK for task {task_id}" + ) + + # Update database + await self._store_task_iteration(task_id, current_iteration) + await self._update_task_status(task_id, TaskStatus.EXECUTING) + + logger.info(f"Human response processed for task {task_id}") + return True + + return False diff --git a/services/orchestrator/task_knowledge_extractor.py b/services/orchestrator/task_knowledge_extractor.py index 409755d..80501ba 100644 --- a/services/orchestrator/task_knowledge_extractor.py +++ b/services/orchestrator/task_knowledge_extractor.py @@ -1,879 +1,879 @@ -""" -Task Knowledge Extractor for FuzeAgent - -This module extracts valuable knowledge from completed tasks and feeds it -into the hierarchical knowledge management system. It analyzes task outcomes, -code changes, conversation patterns, and performance metrics to create -reusable organizational knowledge. -""" - -import asyncio -import json -import logging -import re -from dataclasses import dataclass -from datetime import datetime -from typing import Any, Dict, List, Optional, Tuple - -import asyncpg -from sentence_transformers import SentenceTransformer - -from .knowledge_propagation_engine import KnowledgePropagationEngine, PropagationTrigger -from .organization_rag_manager import ( - ContentType, - KnowledgeCategory, - OrganizationRAGManager, - SourceType, -) -from .team_knowledge_manager import TeamKnowledgeManager - -logger = logging.getLogger(__name__) - - -@dataclass -class TaskKnowledgeExtract: - """Represents extracted knowledge from a task""" - - title: str - content: str - content_type: ContentType - category: KnowledgeCategory - confidence_score: float - tags: List[str] - metadata: Dict[str, Any] - success_indicators: List[str] - failure_patterns: List[str] - - -@dataclass -class ExtractionContext: - """Context for knowledge extraction""" - - task_id: str - agent_id: str - team_id: str - organization_id: str - task_data: Dict[str, Any] - execution_result: Dict[str, Any] - conversation_history: List[Dict[str, Any]] - code_changes: List[Dict[str, Any]] - performance_metrics: Dict[str, Any] - iteration_count: int - total_duration_minutes: float - success: bool - - -class TaskKnowledgeExtractor: - """ - Extracts knowledge from completed tasks and integrates it - into the hierarchical knowledge management system. - """ - - def __init__( - self, - database_url: str, - org_rag_manager: OrganizationRAGManager, - team_knowledge_manager: TeamKnowledgeManager, - propagation_engine: KnowledgePropagationEngine, - ): - self.database_url = database_url - self.org_rag_manager = org_rag_manager - self.team_knowledge_manager = team_knowledge_manager - self.propagation_engine = propagation_engine - self.pool: Optional[asyncpg.Pool] = None - - # Initialize text analysis model - self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") - - # Extraction patterns and rules - self.code_patterns = self._initialize_code_patterns() - self.success_patterns = self._initialize_success_patterns() - self.failure_patterns = self._initialize_failure_patterns() - - # Configuration - self.min_extraction_confidence = 0.4 - self.min_task_duration_minutes = 5 # Don't extract from very short tasks - self.max_content_length = 5000 - - # Statistics - self.extractions_performed = 0 - self.knowledge_items_created = 0 - self.propagations_triggered = 0 - - async def initialize(self): - """Initialize the knowledge extractor""" - logger.info("Initializing TaskKnowledgeExtractor") - - try: - self.pool = await asyncpg.create_pool( - self.database_url, min_size=1, max_size=5, command_timeout=60 - ) - - logger.info("TaskKnowledgeExtractor initialized successfully") - - except Exception as e: - logger.error(f"Failed to initialize TaskKnowledgeExtractor: {e}") - raise - - async def close(self): - """Close database connections""" - if self.pool: - await self.pool.close() - logger.info("TaskKnowledgeExtractor closed") - - async def extract_knowledge_from_task( - self, task_id: str, agent_id: str, execution_result: Dict[str, Any] - ) -> List[str]: - """Extract knowledge from a completed task and store it""" - - try: - # Build extraction context - context = await self._build_extraction_context( - task_id, agent_id, execution_result - ) - - if not context: - logger.warning(f"Could not build extraction context for task {task_id}") - return [] - - # Skip extraction for very short or trivial tasks - if context.total_duration_minutes < self.min_task_duration_minutes: - logger.debug( - f"Skipping extraction for short task {task_id} ({context.total_duration_minutes:.1f}m)" - ) - return [] - - # Extract knowledge items - knowledge_extracts = await self._extract_knowledge_items(context) - - if not knowledge_extracts: - logger.debug(f"No knowledge extracted from task {task_id}") - return [] - - # Store extracted knowledge - stored_knowledge_ids = [] - for extract in knowledge_extracts: - if extract.confidence_score >= self.min_extraction_confidence: - knowledge_id = await self._store_knowledge_extract(context, extract) - if knowledge_id: - stored_knowledge_ids.append(knowledge_id) - - # Trigger knowledge propagation if we have valuable knowledge - if stored_knowledge_ids: - propagation_ids = ( - await self.propagation_engine.trigger_agent_to_team_propagation( - agent_id=context.agent_id, - task_id=context.task_id, - task_outcome={ - "success": context.success, - "task_type": context.task_data.get("task_type", "unknown"), - "complexity": self._assess_task_complexity(context), - "duration_minutes": context.total_duration_minutes, - "knowledge_extracted": len(stored_knowledge_ids), - "iteration_count": context.iteration_count, - }, - ) - ) - - self.propagations_triggered += len(propagation_ids) - logger.info( - f"Triggered {len(propagation_ids)} knowledge propagations for task {task_id}" - ) - - self.extractions_performed += 1 - self.knowledge_items_created += len(stored_knowledge_ids) - - logger.info( - f"Extracted {len(stored_knowledge_ids)} knowledge items from task {task_id}" - ) - return stored_knowledge_ids - - except Exception as e: - logger.error(f"Error extracting knowledge from task {task_id}: {e}") - return [] - - async def _build_extraction_context( - self, task_id: str, agent_id: str, execution_result: Dict[str, Any] - ) -> Optional[ExtractionContext]: - """Build context for knowledge extraction""" - - async with self.pool.acquire() as conn: - # Get basic task information - task_data = await conn.fetchrow( - """ - SELECT t.*, a.team_id, te.organization_id - FROM tasks t - JOIN agents a ON t.agent_id = a.id - JOIN teams te ON a.team_id = te.id - WHERE t.id = $1 - """, - task_id, - ) - - if not task_data: - return None - - # Get conversation history - conversation_history = await conn.fetch( - """ - SELECT message_type, content, metadata, created_at - FROM claude_conversations - WHERE task_id = $1 - ORDER BY created_at ASC - """, - task_id, - ) - - # Get code generations - code_changes = await conn.fetch( - """ - SELECT file_path, file_type, language, content, test_results, quality_metrics - FROM code_generations - WHERE task_id = $1 - ORDER BY generated_at ASC - """, - task_id, - ) - - # Get performance metrics - performance_metrics = await conn.fetch( - """ - SELECT metric_type, metric_value, metric_unit, context - FROM agent_performance_metrics - WHERE task_id = $1 - """, - task_id, - ) - - # Calculate duration - started_at = task_data["started_at"] - completed_at = execution_result.get("completed_at") - if completed_at: - if isinstance(completed_at, str): - completed_at = datetime.fromisoformat( - completed_at.replace("Z", "+00:00") - ) - duration = (completed_at - started_at).total_seconds() / 60.0 - else: - duration = 0.0 - - return ExtractionContext( - task_id=str(task_data["id"]), - agent_id=str(task_data["agent_id"]), - team_id=str(task_data["team_id"]), - organization_id=str(task_data["organization_id"]), - task_data=dict(task_data), - execution_result=execution_result, - conversation_history=[dict(conv) for conv in conversation_history], - code_changes=[dict(code) for code in code_changes], - performance_metrics={ - pm["metric_type"]: pm for pm in performance_metrics - }, - iteration_count=execution_result.get("iterations", 0), - total_duration_minutes=duration, - success=execution_result.get("status") == "completed", - ) - - async def _extract_knowledge_items( - self, context: ExtractionContext - ) -> List[TaskKnowledgeExtract]: - """Extract specific knowledge items from the task context""" - - knowledge_extracts = [] - - # Extract different types of knowledge - knowledge_extracts.extend(await self._extract_code_patterns(context)) - knowledge_extracts.extend(await self._extract_problem_solutions(context)) - knowledge_extracts.extend(await self._extract_debugging_insights(context)) - knowledge_extracts.extend(await self._extract_process_knowledge(context)) - knowledge_extracts.extend(await self._extract_error_patterns(context)) - knowledge_extracts.extend(await self._extract_optimization_insights(context)) - - return knowledge_extracts - - async def _extract_code_patterns( - self, context: ExtractionContext - ) -> List[TaskKnowledgeExtract]: - """Extract reusable code patterns and best practices""" - - extracts = [] - - for code_change in context.code_changes: - if code_change["file_type"] == "implementation": - content = code_change["content"] - language = code_change.get("language", "unknown") - - # Look for reusable patterns - patterns_found = [] - for pattern_name, pattern_info in self.code_patterns.items(): - if any( - keyword in content.lower() - for keyword in pattern_info["keywords"] - ): - patterns_found.append(pattern_name) - - if patterns_found and len(content) > 100: # Substantial code - # Create knowledge extract - title = f"Code Pattern: {', '.join(patterns_found)} ({language})" - extract_content = self._create_code_pattern_content( - content, patterns_found, context - ) - - confidence = self._calculate_code_pattern_confidence( - content, patterns_found, context - ) - - if confidence >= self.min_extraction_confidence: - extract = TaskKnowledgeExtract( - title=title, - content=extract_content, - content_type=ContentType.CODE, - category=KnowledgeCategory.DEVELOPMENT, - confidence_score=confidence, - tags=["code_pattern", language, *patterns_found], - metadata={ - "language": language, - "file_path": code_change["file_path"], - "patterns": patterns_found, - "task_success": context.success, - "lines_of_code": len(content.split("\n")), - }, - success_indicators=self._extract_success_indicators( - context - ), - failure_patterns=[], - ) - - extracts.append(extract) - - return extracts - - async def _extract_problem_solutions( - self, context: ExtractionContext - ) -> List[TaskKnowledgeExtract]: - """Extract problem-solution pairs from the task""" - - extracts = [] - - # Analyze conversation for problem descriptions and solutions - problem_solution_pairs = self._identify_problem_solution_pairs( - context.conversation_history - ) - - for problem, solution in problem_solution_pairs: - if len(problem) > 50 and len(solution) > 50: # Substantial content - title = f"Solution: {problem[:50]}..." - content = f"**Problem:**\n{problem}\n\n**Solution:**\n{solution}" - - # Determine category based on content - category = self._categorize_problem_solution(problem, solution) - - confidence = self._calculate_solution_confidence( - problem, solution, context - ) - - if confidence >= self.min_extraction_confidence: - extract = TaskKnowledgeExtract( - title=title, - content=content[: self.max_content_length], - content_type=ContentType.PROCEDURE, - category=category, - confidence_score=confidence, - tags=["problem_solution", "troubleshooting"], - metadata={ - "problem_type": self._classify_problem_type(problem), - "solution_type": self._classify_solution_type(solution), - "task_success": context.success, - }, - success_indicators=self._extract_success_indicators(context), - failure_patterns=[], - ) - - extracts.append(extract) - - return extracts - - async def _extract_debugging_insights( - self, context: ExtractionContext - ) -> List[TaskKnowledgeExtract]: - """Extract debugging approaches and insights""" - - extracts = [] - - # Look for error messages and resolution patterns - debugging_sessions = self._identify_debugging_sessions( - context.conversation_history - ) - - for session in debugging_sessions: - if session["resolution"] and context.success: - title = f"Debugging: {session['error_type']}" - content = self._create_debugging_content(session) - - confidence = self._calculate_debugging_confidence(session, context) - - if confidence >= self.min_extraction_confidence: - extract = TaskKnowledgeExtract( - title=title, - content=content, - content_type=ContentType.PROCEDURE, - category=KnowledgeCategory.TROUBLESHOOTING, - confidence_score=confidence, - tags=["debugging", session["error_type"], "troubleshooting"], - metadata={ - "error_type": session["error_type"], - "resolution_method": session["resolution_method"], - "tools_used": session.get("tools_used", []), - }, - success_indicators=self._extract_success_indicators(context), - failure_patterns=session.get("failure_patterns", []), - ) - - extracts.append(extract) - - return extracts - - async def _extract_process_knowledge( - self, context: ExtractionContext - ) -> List[TaskKnowledgeExtract]: - """Extract process and workflow knowledge""" - - extracts = [] - - if context.iteration_count > 1: # Multi-iteration tasks have process insights - title = f"Process: {context.task_data.get('task_type', 'Task')} Workflow" - - process_content = self._create_process_content(context) - confidence = self._calculate_process_confidence(context) - - if confidence >= self.min_extraction_confidence: - extract = TaskKnowledgeExtract( - title=title, - content=process_content, - content_type=ContentType.PROCEDURE, - category=KnowledgeCategory.PROCESS, - confidence_score=confidence, - tags=[ - "process", - "workflow", - context.task_data.get("task_type", "general"), - ], - metadata={ - "iterations_used": context.iteration_count, - "duration_minutes": context.total_duration_minutes, - "success_rate": 1.0 if context.success else 0.0, - "complexity": self._assess_task_complexity(context), - }, - success_indicators=self._extract_success_indicators(context), - failure_patterns=[], - ) - - extracts.append(extract) - - return extracts - - async def _extract_error_patterns( - self, context: ExtractionContext - ) -> List[TaskKnowledgeExtract]: - """Extract error patterns and avoidance strategies""" - - extracts = [] - - # Look for error patterns in failed tasks or recovered errors - error_patterns = self._identify_error_patterns(context.conversation_history) - - for pattern in error_patterns: - if pattern["frequency"] >= 2 or pattern["severity"] == "high": - title = f"Error Pattern: {pattern['error_type']}" - content = self._create_error_pattern_content(pattern, context) - - confidence = self._calculate_error_pattern_confidence(pattern, context) - - if confidence >= self.min_extraction_confidence: - extract = TaskKnowledgeExtract( - title=title, - content=content, - content_type=ContentType.DOCUMENTATION, - category=KnowledgeCategory.TROUBLESHOOTING, - confidence_score=confidence, - tags=["error_pattern", pattern["error_type"], "prevention"], - metadata={ - "error_type": pattern["error_type"], - "frequency": pattern["frequency"], - "severity": pattern["severity"], - "prevention_strategies": pattern.get("prevention", []), - }, - success_indicators=[], - failure_patterns=pattern.get("indicators", []), - ) - - extracts.append(extract) - - return extracts - - async def _extract_optimization_insights( - self, context: ExtractionContext - ) -> List[TaskKnowledgeExtract]: - """Extract performance optimization insights""" - - extracts = [] - - # Look for performance improvements in metrics - if "execution_time_minutes" in context.performance_metrics: - perf_data = context.performance_metrics["execution_time_minutes"] - if ( - perf_data["metric_value"] < 30 and context.success - ): # Efficient completion - title = "Performance Optimization: Efficient Task Execution" - content = self._create_optimization_content(context) - - confidence = self._calculate_optimization_confidence(context) - - if confidence >= self.min_extraction_confidence: - extract = TaskKnowledgeExtract( - title=title, - content=content, - content_type=ContentType.BEST_PRACTICE, - category=KnowledgeCategory.DEVELOPMENT, - confidence_score=confidence, - tags=["optimization", "performance", "efficiency"], - metadata={ - "execution_time": perf_data["metric_value"], - "iteration_efficiency": context.iteration_count - / context.total_duration_minutes, - "optimization_techniques": self._identify_optimization_techniques( - context - ), - }, - success_indicators=self._extract_success_indicators(context), - failure_patterns=[], - ) - - extracts.append(extract) - - return extracts - - async def _store_knowledge_extract( - self, context: ExtractionContext, extract: TaskKnowledgeExtract - ) -> Optional[str]: - """Store a knowledge extract in the appropriate knowledge base""" - - try: - # Store in organization knowledge base - knowledge_id = await self.org_rag_manager.add_knowledge( - organization_id=context.organization_id, - title=extract.title, - content=extract.content, - content_type=extract.content_type, - knowledge_category=extract.category, - source_type=SourceType.TASK_OUTCOME, - source_agent_id=context.agent_id, - source_team_id=context.team_id, - source_task_id=context.task_id, - relevance_score=extract.confidence_score, - quality_score=extract.confidence_score, - metadata={ - **extract.metadata, - "extraction_timestamp": datetime.now().isoformat(), - "extractor_version": "1.0", - "success_indicators": extract.success_indicators, - "failure_patterns": extract.failure_patterns, - }, - tags=extract.tags, - ) - - return knowledge_id - - except Exception as e: - logger.error(f"Error storing knowledge extract: {e}") - return None - - # Helper methods for pattern matching and analysis - def _initialize_code_patterns(self) -> Dict[str, Dict[str, Any]]: - """Initialize code pattern definitions""" - return { - "api_integration": { - "keywords": ["fetch", "axios", "request", "api", "endpoint", "rest"], - "confidence_boost": 0.2, - }, - "database_operations": { - "keywords": [ - "select", - "insert", - "update", - "delete", - "query", - "database", - "sql", - ], - "confidence_boost": 0.2, - }, - "authentication": { - "keywords": ["auth", "login", "token", "jwt", "session", "passport"], - "confidence_boost": 0.15, - }, - "error_handling": { - "keywords": ["try", "catch", "error", "exception", "throw"], - "confidence_boost": 0.1, - }, - "testing": { - "keywords": ["test", "spec", "describe", "it", "expect", "mock"], - "confidence_boost": 0.15, - }, - "optimization": { - "keywords": ["performance", "optimize", "cache", "memory", "speed"], - "confidence_boost": 0.2, - }, - } - - def _initialize_success_patterns(self) -> List[str]: - """Initialize success indicator patterns""" - return [ - r"test.*pass", - r"build.*success", - r"deploy.*complete", - r"fix.*issue", - r"resolve.*problem", - r"implement.*feature", - r"complete.*task", - ] - - def _initialize_failure_patterns(self) -> List[str]: - """Initialize failure indicator patterns""" - return [ - r"error.*occur", - r"fail.*to", - r"timeout.*exceed", - r"connection.*refuse", - r"not.*found", - r"access.*deni", - r"invalid.*request", - ] - - def _assess_task_complexity(self, context: ExtractionContext) -> str: - """Assess task complexity based on various factors""" - - complexity_score = 0 - - # Factor 1: Iteration count - if context.iteration_count > 10: - complexity_score += 3 - elif context.iteration_count > 5: - complexity_score += 2 - elif context.iteration_count > 2: - complexity_score += 1 - - # Factor 2: Duration - if context.total_duration_minutes > 180: # 3 hours - complexity_score += 3 - elif context.total_duration_minutes > 60: # 1 hour - complexity_score += 2 - elif context.total_duration_minutes > 30: - complexity_score += 1 - - # Factor 3: Code changes - if len(context.code_changes) > 10: - complexity_score += 2 - elif len(context.code_changes) > 5: - complexity_score += 1 - - # Factor 4: Conversation length - if len(context.conversation_history) > 50: - complexity_score += 2 - elif len(context.conversation_history) > 20: - complexity_score += 1 - - if complexity_score >= 6: - return "very_high" - elif complexity_score >= 4: - return "high" - elif complexity_score >= 2: - return "medium" - else: - return "low" - - def _extract_success_indicators(self, context: ExtractionContext) -> List[str]: - """Extract success indicators from the task execution""" - - indicators = [] - - # Look for success patterns in conversation - for conv in context.conversation_history: - content = conv.get("content", "").lower() - for pattern in self.success_patterns: - if re.search(pattern, content): - indicators.append(pattern) - - # Add task-specific indicators - if context.success: - indicators.append("task_completed_successfully") - - if context.execution_result.get("pull_request_url"): - indicators.append("pull_request_created") - - return list(set(indicators)) # Remove duplicates - - # Additional helper methods would be implemented here... - # (The file is getting quite long, so I'll implement key methods and indicate where others would go) - - def _identify_problem_solution_pairs( - self, conversation_history: List[Dict] - ) -> List[Tuple[str, str]]: - """Identify problem-solution pairs in conversation history""" - pairs = [] - # Implementation would analyze conversation flow to identify problems and their solutions - # This is a simplified placeholder - return pairs - - def _categorize_problem_solution( - self, problem: str, solution: str - ) -> KnowledgeCategory: - """Categorize a problem-solution pair""" - # Simple categorization based on keywords - combined_text = (problem + " " + solution).lower() - - if any(word in combined_text for word in ["test", "testing", "spec"]): - return KnowledgeCategory.TESTING - elif any(word in combined_text for word in ["deploy", "build", "ci", "cd"]): - return KnowledgeCategory.INFRASTRUCTURE - elif any(word in combined_text for word in ["security", "auth", "permission"]): - return KnowledgeCategory.SECURITY - elif any(word in combined_text for word in ["design", "ui", "ux", "interface"]): - return KnowledgeCategory.DESIGN - else: - return KnowledgeCategory.DEVELOPMENT - - def _calculate_code_pattern_confidence( - self, content: str, patterns: List[str], context: ExtractionContext - ) -> float: - """Calculate confidence score for code pattern extraction""" - base_confidence = 0.5 - - # Boost for successful task - if context.success: - base_confidence += 0.2 - - # Boost for multiple patterns - if len(patterns) > 1: - base_confidence += 0.1 - - # Boost for substantial code - if len(content) > 500: - base_confidence += 0.1 - - return min(1.0, base_confidence) - - def _calculate_solution_confidence( - self, problem: str, solution: str, context: ExtractionContext - ) -> float: - """Calculate confidence score for solution extraction""" - base_confidence = 0.4 - - if context.success: - base_confidence += 0.3 - - if len(solution) > 200: # Detailed solution - base_confidence += 0.1 - - return min(1.0, base_confidence) - - def _calculate_debugging_confidence( - self, session: Dict, context: ExtractionContext - ) -> float: - """Calculate confidence for debugging insights""" - base_confidence = 0.6 if context.success else 0.3 - - if session.get("resolution_method"): - base_confidence += 0.2 - - return min(1.0, base_confidence) - - def _calculate_process_confidence(self, context: ExtractionContext) -> float: - """Calculate confidence for process knowledge""" - if not context.success: - return 0.2 - - # Base confidence increases with iteration count (more process learning) - base_confidence = min(0.8, 0.3 + (context.iteration_count * 0.05)) - - return base_confidence - - def _calculate_error_pattern_confidence( - self, pattern: Dict, context: ExtractionContext - ) -> float: - """Calculate confidence for error pattern extraction""" - base_confidence = 0.4 - - if pattern["frequency"] > 2: - base_confidence += 0.2 - - if pattern["severity"] == "high": - base_confidence += 0.2 - - return min(1.0, base_confidence) - - def _calculate_optimization_confidence(self, context: ExtractionContext) -> float: - """Calculate confidence for optimization insights""" - if not context.success: - return 0.1 - - base_confidence = 0.5 - - # Boost for efficient execution - if context.total_duration_minutes < 30: - base_confidence += 0.2 - - if context.iteration_count < 5: - base_confidence += 0.1 - - return min(1.0, base_confidence) - - # Content creation methods (simplified implementations) - def _create_code_pattern_content( - self, content: str, patterns: List[str], context: ExtractionContext - ) -> str: - """Create formatted content for code pattern knowledge""" - return f"**Code Pattern: {', '.join(patterns)}**\n\n{content[:2000]}..." - - def _create_debugging_content(self, session: Dict) -> str: - """Create formatted content for debugging knowledge""" - return f"**Error:** {session.get('error_type', 'Unknown')}\n\n**Resolution:** {session.get('resolution', 'No resolution provided')}" - - def _create_process_content(self, context: ExtractionContext) -> str: - """Create formatted content for process knowledge""" - return f"**Task Type:** {context.task_data.get('task_type', 'Unknown')}\n**Iterations:** {context.iteration_count}\n**Duration:** {context.total_duration_minutes:.1f} minutes\n**Success:** {'Yes' if context.success else 'No'}" - - def _create_error_pattern_content( - self, pattern: Dict, context: ExtractionContext - ) -> str: - """Create formatted content for error pattern knowledge""" - return f"**Error Type:** {pattern['error_type']}\n**Frequency:** {pattern['frequency']}\n**Prevention:** {', '.join(pattern.get('prevention', []))}" - - def _create_optimization_content(self, context: ExtractionContext) -> str: - """Create formatted content for optimization knowledge""" - return f"**Optimization achieved in {context.total_duration_minutes:.1f} minutes with {context.iteration_count} iterations**" - - # Placeholder methods for more complex analysis functions - def _identify_debugging_sessions( - self, conversation_history: List[Dict] - ) -> List[Dict]: - """Identify debugging sessions in conversation history""" - return [] # Simplified implementation - - def _identify_error_patterns(self, conversation_history: List[Dict]) -> List[Dict]: - """Identify error patterns in conversation history""" - return [] # Simplified implementation - - def _identify_optimization_techniques( - self, context: ExtractionContext - ) -> List[str]: - """Identify optimization techniques used""" - return [] # Simplified implementation - - def _classify_problem_type(self, problem: str) -> str: - """Classify the type of problem""" - return "general" # Simplified implementation - - def _classify_solution_type(self, solution: str) -> str: - """Classify the type of solution""" - return "general" # Simplified implementation +""" +Task Knowledge Extractor for FuzeAgent + +This module extracts valuable knowledge from completed tasks and feeds it +into the hierarchical knowledge management system. It analyzes task outcomes, +code changes, conversation patterns, and performance metrics to create +reusable organizational knowledge. +""" + +import asyncio +import json +import logging +import re +from dataclasses import dataclass +from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple + +import asyncpg +from sentence_transformers import SentenceTransformer + +from .knowledge_propagation_engine import KnowledgePropagationEngine, PropagationTrigger +from .organization_rag_manager import ( + ContentType, + KnowledgeCategory, + OrganizationRAGManager, + SourceType, +) +from .team_knowledge_manager import TeamKnowledgeManager + +logger = logging.getLogger(__name__) + + +@dataclass +class TaskKnowledgeExtract: + """Represents extracted knowledge from a task""" + + title: str + content: str + content_type: ContentType + category: KnowledgeCategory + confidence_score: float + tags: List[str] + metadata: Dict[str, Any] + success_indicators: List[str] + failure_patterns: List[str] + + +@dataclass +class ExtractionContext: + """Context for knowledge extraction""" + + task_id: str + agent_id: str + team_id: str + organization_id: str + task_data: Dict[str, Any] + execution_result: Dict[str, Any] + conversation_history: List[Dict[str, Any]] + code_changes: List[Dict[str, Any]] + performance_metrics: Dict[str, Any] + iteration_count: int + total_duration_minutes: float + success: bool + + +class TaskKnowledgeExtractor: + """ + Extracts knowledge from completed tasks and integrates it + into the hierarchical knowledge management system. + """ + + def __init__( + self, + database_url: str, + org_rag_manager: OrganizationRAGManager, + team_knowledge_manager: TeamKnowledgeManager, + propagation_engine: KnowledgePropagationEngine, + ): + self.database_url = database_url + self.org_rag_manager = org_rag_manager + self.team_knowledge_manager = team_knowledge_manager + self.propagation_engine = propagation_engine + self.pool: Optional[asyncpg.Pool] = None + + # Initialize text analysis model + self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") + + # Extraction patterns and rules + self.code_patterns = self._initialize_code_patterns() + self.success_patterns = self._initialize_success_patterns() + self.failure_patterns = self._initialize_failure_patterns() + + # Configuration + self.min_extraction_confidence = 0.4 + self.min_task_duration_minutes = 5 # Don't extract from very short tasks + self.max_content_length = 5000 + + # Statistics + self.extractions_performed = 0 + self.knowledge_items_created = 0 + self.propagations_triggered = 0 + + async def initialize(self): + """Initialize the knowledge extractor""" + logger.info("Initializing TaskKnowledgeExtractor") + + try: + self.pool = await asyncpg.create_pool( + self.database_url, min_size=1, max_size=5, command_timeout=60 + ) + + logger.info("TaskKnowledgeExtractor initialized successfully") + + except Exception as e: + logger.error(f"Failed to initialize TaskKnowledgeExtractor: {e}") + raise + + async def close(self): + """Close database connections""" + if self.pool: + await self.pool.close() + logger.info("TaskKnowledgeExtractor closed") + + async def extract_knowledge_from_task( + self, task_id: str, agent_id: str, execution_result: Dict[str, Any] + ) -> List[str]: + """Extract knowledge from a completed task and store it""" + + try: + # Build extraction context + context = await self._build_extraction_context( + task_id, agent_id, execution_result + ) + + if not context: + logger.warning(f"Could not build extraction context for task {task_id}") + return [] + + # Skip extraction for very short or trivial tasks + if context.total_duration_minutes < self.min_task_duration_minutes: + logger.debug( + f"Skipping extraction for short task {task_id} ({context.total_duration_minutes:.1f}m)" + ) + return [] + + # Extract knowledge items + knowledge_extracts = await self._extract_knowledge_items(context) + + if not knowledge_extracts: + logger.debug(f"No knowledge extracted from task {task_id}") + return [] + + # Store extracted knowledge + stored_knowledge_ids = [] + for extract in knowledge_extracts: + if extract.confidence_score >= self.min_extraction_confidence: + knowledge_id = await self._store_knowledge_extract(context, extract) + if knowledge_id: + stored_knowledge_ids.append(knowledge_id) + + # Trigger knowledge propagation if we have valuable knowledge + if stored_knowledge_ids: + propagation_ids = ( + await self.propagation_engine.trigger_agent_to_team_propagation( + agent_id=context.agent_id, + task_id=context.task_id, + task_outcome={ + "success": context.success, + "task_type": context.task_data.get("task_type", "unknown"), + "complexity": self._assess_task_complexity(context), + "duration_minutes": context.total_duration_minutes, + "knowledge_extracted": len(stored_knowledge_ids), + "iteration_count": context.iteration_count, + }, + ) + ) + + self.propagations_triggered += len(propagation_ids) + logger.info( + f"Triggered {len(propagation_ids)} knowledge propagations for task {task_id}" + ) + + self.extractions_performed += 1 + self.knowledge_items_created += len(stored_knowledge_ids) + + logger.info( + f"Extracted {len(stored_knowledge_ids)} knowledge items from task {task_id}" + ) + return stored_knowledge_ids + + except Exception as e: + logger.error(f"Error extracting knowledge from task {task_id}: {e}") + return [] + + async def _build_extraction_context( + self, task_id: str, agent_id: str, execution_result: Dict[str, Any] + ) -> Optional[ExtractionContext]: + """Build context for knowledge extraction""" + + async with self.pool.acquire() as conn: + # Get basic task information + task_data = await conn.fetchrow( + """ + SELECT t.*, a.team_id, te.organization_id + FROM tasks t + JOIN agents a ON t.agent_id = a.id + JOIN teams te ON a.team_id = te.id + WHERE t.id = $1 + """, + task_id, + ) + + if not task_data: + return None + + # Get conversation history + conversation_history = await conn.fetch( + """ + SELECT message_type, content, metadata, created_at + FROM claude_conversations + WHERE task_id = $1 + ORDER BY created_at ASC + """, + task_id, + ) + + # Get code generations + code_changes = await conn.fetch( + """ + SELECT file_path, file_type, language, content, test_results, quality_metrics + FROM code_generations + WHERE task_id = $1 + ORDER BY generated_at ASC + """, + task_id, + ) + + # Get performance metrics + performance_metrics = await conn.fetch( + """ + SELECT metric_type, metric_value, metric_unit, context + FROM agent_performance_metrics + WHERE task_id = $1 + """, + task_id, + ) + + # Calculate duration + started_at = task_data["started_at"] + completed_at = execution_result.get("completed_at") + if completed_at: + if isinstance(completed_at, str): + completed_at = datetime.fromisoformat( + completed_at.replace("Z", "+00:00") + ) + duration = (completed_at - started_at).total_seconds() / 60.0 + else: + duration = 0.0 + + return ExtractionContext( + task_id=str(task_data["id"]), + agent_id=str(task_data["agent_id"]), + team_id=str(task_data["team_id"]), + organization_id=str(task_data["organization_id"]), + task_data=dict(task_data), + execution_result=execution_result, + conversation_history=[dict(conv) for conv in conversation_history], + code_changes=[dict(code) for code in code_changes], + performance_metrics={ + pm["metric_type"]: pm for pm in performance_metrics + }, + iteration_count=execution_result.get("iterations", 0), + total_duration_minutes=duration, + success=execution_result.get("status") == "completed", + ) + + async def _extract_knowledge_items( + self, context: ExtractionContext + ) -> List[TaskKnowledgeExtract]: + """Extract specific knowledge items from the task context""" + + knowledge_extracts = [] + + # Extract different types of knowledge + knowledge_extracts.extend(await self._extract_code_patterns(context)) + knowledge_extracts.extend(await self._extract_problem_solutions(context)) + knowledge_extracts.extend(await self._extract_debugging_insights(context)) + knowledge_extracts.extend(await self._extract_process_knowledge(context)) + knowledge_extracts.extend(await self._extract_error_patterns(context)) + knowledge_extracts.extend(await self._extract_optimization_insights(context)) + + return knowledge_extracts + + async def _extract_code_patterns( + self, context: ExtractionContext + ) -> List[TaskKnowledgeExtract]: + """Extract reusable code patterns and best practices""" + + extracts = [] + + for code_change in context.code_changes: + if code_change["file_type"] == "implementation": + content = code_change["content"] + language = code_change.get("language", "unknown") + + # Look for reusable patterns + patterns_found = [] + for pattern_name, pattern_info in self.code_patterns.items(): + if any( + keyword in content.lower() + for keyword in pattern_info["keywords"] + ): + patterns_found.append(pattern_name) + + if patterns_found and len(content) > 100: # Substantial code + # Create knowledge extract + title = f"Code Pattern: {', '.join(patterns_found)} ({language})" + extract_content = self._create_code_pattern_content( + content, patterns_found, context + ) + + confidence = self._calculate_code_pattern_confidence( + content, patterns_found, context + ) + + if confidence >= self.min_extraction_confidence: + extract = TaskKnowledgeExtract( + title=title, + content=extract_content, + content_type=ContentType.CODE, + category=KnowledgeCategory.DEVELOPMENT, + confidence_score=confidence, + tags=["code_pattern", language, *patterns_found], + metadata={ + "language": language, + "file_path": code_change["file_path"], + "patterns": patterns_found, + "task_success": context.success, + "lines_of_code": len(content.split("\n")), + }, + success_indicators=self._extract_success_indicators( + context + ), + failure_patterns=[], + ) + + extracts.append(extract) + + return extracts + + async def _extract_problem_solutions( + self, context: ExtractionContext + ) -> List[TaskKnowledgeExtract]: + """Extract problem-solution pairs from the task""" + + extracts = [] + + # Analyze conversation for problem descriptions and solutions + problem_solution_pairs = self._identify_problem_solution_pairs( + context.conversation_history + ) + + for problem, solution in problem_solution_pairs: + if len(problem) > 50 and len(solution) > 50: # Substantial content + title = f"Solution: {problem[:50]}..." + content = f"**Problem:**\n{problem}\n\n**Solution:**\n{solution}" + + # Determine category based on content + category = self._categorize_problem_solution(problem, solution) + + confidence = self._calculate_solution_confidence( + problem, solution, context + ) + + if confidence >= self.min_extraction_confidence: + extract = TaskKnowledgeExtract( + title=title, + content=content[: self.max_content_length], + content_type=ContentType.PROCEDURE, + category=category, + confidence_score=confidence, + tags=["problem_solution", "troubleshooting"], + metadata={ + "problem_type": self._classify_problem_type(problem), + "solution_type": self._classify_solution_type(solution), + "task_success": context.success, + }, + success_indicators=self._extract_success_indicators(context), + failure_patterns=[], + ) + + extracts.append(extract) + + return extracts + + async def _extract_debugging_insights( + self, context: ExtractionContext + ) -> List[TaskKnowledgeExtract]: + """Extract debugging approaches and insights""" + + extracts = [] + + # Look for error messages and resolution patterns + debugging_sessions = self._identify_debugging_sessions( + context.conversation_history + ) + + for session in debugging_sessions: + if session["resolution"] and context.success: + title = f"Debugging: {session['error_type']}" + content = self._create_debugging_content(session) + + confidence = self._calculate_debugging_confidence(session, context) + + if confidence >= self.min_extraction_confidence: + extract = TaskKnowledgeExtract( + title=title, + content=content, + content_type=ContentType.PROCEDURE, + category=KnowledgeCategory.TROUBLESHOOTING, + confidence_score=confidence, + tags=["debugging", session["error_type"], "troubleshooting"], + metadata={ + "error_type": session["error_type"], + "resolution_method": session["resolution_method"], + "tools_used": session.get("tools_used", []), + }, + success_indicators=self._extract_success_indicators(context), + failure_patterns=session.get("failure_patterns", []), + ) + + extracts.append(extract) + + return extracts + + async def _extract_process_knowledge( + self, context: ExtractionContext + ) -> List[TaskKnowledgeExtract]: + """Extract process and workflow knowledge""" + + extracts = [] + + if context.iteration_count > 1: # Multi-iteration tasks have process insights + title = f"Process: {context.task_data.get('task_type', 'Task')} Workflow" + + process_content = self._create_process_content(context) + confidence = self._calculate_process_confidence(context) + + if confidence >= self.min_extraction_confidence: + extract = TaskKnowledgeExtract( + title=title, + content=process_content, + content_type=ContentType.PROCEDURE, + category=KnowledgeCategory.PROCESS, + confidence_score=confidence, + tags=[ + "process", + "workflow", + context.task_data.get("task_type", "general"), + ], + metadata={ + "iterations_used": context.iteration_count, + "duration_minutes": context.total_duration_minutes, + "success_rate": 1.0 if context.success else 0.0, + "complexity": self._assess_task_complexity(context), + }, + success_indicators=self._extract_success_indicators(context), + failure_patterns=[], + ) + + extracts.append(extract) + + return extracts + + async def _extract_error_patterns( + self, context: ExtractionContext + ) -> List[TaskKnowledgeExtract]: + """Extract error patterns and avoidance strategies""" + + extracts = [] + + # Look for error patterns in failed tasks or recovered errors + error_patterns = self._identify_error_patterns(context.conversation_history) + + for pattern in error_patterns: + if pattern["frequency"] >= 2 or pattern["severity"] == "high": + title = f"Error Pattern: {pattern['error_type']}" + content = self._create_error_pattern_content(pattern, context) + + confidence = self._calculate_error_pattern_confidence(pattern, context) + + if confidence >= self.min_extraction_confidence: + extract = TaskKnowledgeExtract( + title=title, + content=content, + content_type=ContentType.DOCUMENTATION, + category=KnowledgeCategory.TROUBLESHOOTING, + confidence_score=confidence, + tags=["error_pattern", pattern["error_type"], "prevention"], + metadata={ + "error_type": pattern["error_type"], + "frequency": pattern["frequency"], + "severity": pattern["severity"], + "prevention_strategies": pattern.get("prevention", []), + }, + success_indicators=[], + failure_patterns=pattern.get("indicators", []), + ) + + extracts.append(extract) + + return extracts + + async def _extract_optimization_insights( + self, context: ExtractionContext + ) -> List[TaskKnowledgeExtract]: + """Extract performance optimization insights""" + + extracts = [] + + # Look for performance improvements in metrics + if "execution_time_minutes" in context.performance_metrics: + perf_data = context.performance_metrics["execution_time_minutes"] + if ( + perf_data["metric_value"] < 30 and context.success + ): # Efficient completion + title = "Performance Optimization: Efficient Task Execution" + content = self._create_optimization_content(context) + + confidence = self._calculate_optimization_confidence(context) + + if confidence >= self.min_extraction_confidence: + extract = TaskKnowledgeExtract( + title=title, + content=content, + content_type=ContentType.BEST_PRACTICE, + category=KnowledgeCategory.DEVELOPMENT, + confidence_score=confidence, + tags=["optimization", "performance", "efficiency"], + metadata={ + "execution_time": perf_data["metric_value"], + "iteration_efficiency": context.iteration_count + / context.total_duration_minutes, + "optimization_techniques": self._identify_optimization_techniques( + context + ), + }, + success_indicators=self._extract_success_indicators(context), + failure_patterns=[], + ) + + extracts.append(extract) + + return extracts + + async def _store_knowledge_extract( + self, context: ExtractionContext, extract: TaskKnowledgeExtract + ) -> Optional[str]: + """Store a knowledge extract in the appropriate knowledge base""" + + try: + # Store in organization knowledge base + knowledge_id = await self.org_rag_manager.add_knowledge( + organization_id=context.organization_id, + title=extract.title, + content=extract.content, + content_type=extract.content_type, + knowledge_category=extract.category, + source_type=SourceType.TASK_OUTCOME, + source_agent_id=context.agent_id, + source_team_id=context.team_id, + source_task_id=context.task_id, + relevance_score=extract.confidence_score, + quality_score=extract.confidence_score, + metadata={ + **extract.metadata, + "extraction_timestamp": datetime.now().isoformat(), + "extractor_version": "1.0", + "success_indicators": extract.success_indicators, + "failure_patterns": extract.failure_patterns, + }, + tags=extract.tags, + ) + + return knowledge_id + + except Exception as e: + logger.error(f"Error storing knowledge extract: {e}") + return None + + # Helper methods for pattern matching and analysis + def _initialize_code_patterns(self) -> Dict[str, Dict[str, Any]]: + """Initialize code pattern definitions""" + return { + "api_integration": { + "keywords": ["fetch", "axios", "request", "api", "endpoint", "rest"], + "confidence_boost": 0.2, + }, + "database_operations": { + "keywords": [ + "select", + "insert", + "update", + "delete", + "query", + "database", + "sql", + ], + "confidence_boost": 0.2, + }, + "authentication": { + "keywords": ["auth", "login", "token", "jwt", "session", "passport"], + "confidence_boost": 0.15, + }, + "error_handling": { + "keywords": ["try", "catch", "error", "exception", "throw"], + "confidence_boost": 0.1, + }, + "testing": { + "keywords": ["test", "spec", "describe", "it", "expect", "mock"], + "confidence_boost": 0.15, + }, + "optimization": { + "keywords": ["performance", "optimize", "cache", "memory", "speed"], + "confidence_boost": 0.2, + }, + } + + def _initialize_success_patterns(self) -> List[str]: + """Initialize success indicator patterns""" + return [ + r"test.*pass", + r"build.*success", + r"deploy.*complete", + r"fix.*issue", + r"resolve.*problem", + r"implement.*feature", + r"complete.*task", + ] + + def _initialize_failure_patterns(self) -> List[str]: + """Initialize failure indicator patterns""" + return [ + r"error.*occur", + r"fail.*to", + r"timeout.*exceed", + r"connection.*refuse", + r"not.*found", + r"access.*deni", + r"invalid.*request", + ] + + def _assess_task_complexity(self, context: ExtractionContext) -> str: + """Assess task complexity based on various factors""" + + complexity_score = 0 + + # Factor 1: Iteration count + if context.iteration_count > 10: + complexity_score += 3 + elif context.iteration_count > 5: + complexity_score += 2 + elif context.iteration_count > 2: + complexity_score += 1 + + # Factor 2: Duration + if context.total_duration_minutes > 180: # 3 hours + complexity_score += 3 + elif context.total_duration_minutes > 60: # 1 hour + complexity_score += 2 + elif context.total_duration_minutes > 30: + complexity_score += 1 + + # Factor 3: Code changes + if len(context.code_changes) > 10: + complexity_score += 2 + elif len(context.code_changes) > 5: + complexity_score += 1 + + # Factor 4: Conversation length + if len(context.conversation_history) > 50: + complexity_score += 2 + elif len(context.conversation_history) > 20: + complexity_score += 1 + + if complexity_score >= 6: + return "very_high" + elif complexity_score >= 4: + return "high" + elif complexity_score >= 2: + return "medium" + else: + return "low" + + def _extract_success_indicators(self, context: ExtractionContext) -> List[str]: + """Extract success indicators from the task execution""" + + indicators = [] + + # Look for success patterns in conversation + for conv in context.conversation_history: + content = conv.get("content", "").lower() + for pattern in self.success_patterns: + if re.search(pattern, content): + indicators.append(pattern) + + # Add task-specific indicators + if context.success: + indicators.append("task_completed_successfully") + + if context.execution_result.get("pull_request_url"): + indicators.append("pull_request_created") + + return list(set(indicators)) # Remove duplicates + + # Additional helper methods would be implemented here... + # (The file is getting quite long, so I'll implement key methods and indicate where others would go) + + def _identify_problem_solution_pairs( + self, conversation_history: List[Dict] + ) -> List[Tuple[str, str]]: + """Identify problem-solution pairs in conversation history""" + pairs = [] + # Implementation would analyze conversation flow to identify problems and their solutions + # This is a simplified placeholder + return pairs + + def _categorize_problem_solution( + self, problem: str, solution: str + ) -> KnowledgeCategory: + """Categorize a problem-solution pair""" + # Simple categorization based on keywords + combined_text = (problem + " " + solution).lower() + + if any(word in combined_text for word in ["test", "testing", "spec"]): + return KnowledgeCategory.TESTING + elif any(word in combined_text for word in ["deploy", "build", "ci", "cd"]): + return KnowledgeCategory.INFRASTRUCTURE + elif any(word in combined_text for word in ["security", "auth", "permission"]): + return KnowledgeCategory.SECURITY + elif any(word in combined_text for word in ["design", "ui", "ux", "interface"]): + return KnowledgeCategory.DESIGN + else: + return KnowledgeCategory.DEVELOPMENT + + def _calculate_code_pattern_confidence( + self, content: str, patterns: List[str], context: ExtractionContext + ) -> float: + """Calculate confidence score for code pattern extraction""" + base_confidence = 0.5 + + # Boost for successful task + if context.success: + base_confidence += 0.2 + + # Boost for multiple patterns + if len(patterns) > 1: + base_confidence += 0.1 + + # Boost for substantial code + if len(content) > 500: + base_confidence += 0.1 + + return min(1.0, base_confidence) + + def _calculate_solution_confidence( + self, problem: str, solution: str, context: ExtractionContext + ) -> float: + """Calculate confidence score for solution extraction""" + base_confidence = 0.4 + + if context.success: + base_confidence += 0.3 + + if len(solution) > 200: # Detailed solution + base_confidence += 0.1 + + return min(1.0, base_confidence) + + def _calculate_debugging_confidence( + self, session: Dict, context: ExtractionContext + ) -> float: + """Calculate confidence for debugging insights""" + base_confidence = 0.6 if context.success else 0.3 + + if session.get("resolution_method"): + base_confidence += 0.2 + + return min(1.0, base_confidence) + + def _calculate_process_confidence(self, context: ExtractionContext) -> float: + """Calculate confidence for process knowledge""" + if not context.success: + return 0.2 + + # Base confidence increases with iteration count (more process learning) + base_confidence = min(0.8, 0.3 + (context.iteration_count * 0.05)) + + return base_confidence + + def _calculate_error_pattern_confidence( + self, pattern: Dict, context: ExtractionContext + ) -> float: + """Calculate confidence for error pattern extraction""" + base_confidence = 0.4 + + if pattern["frequency"] > 2: + base_confidence += 0.2 + + if pattern["severity"] == "high": + base_confidence += 0.2 + + return min(1.0, base_confidence) + + def _calculate_optimization_confidence(self, context: ExtractionContext) -> float: + """Calculate confidence for optimization insights""" + if not context.success: + return 0.1 + + base_confidence = 0.5 + + # Boost for efficient execution + if context.total_duration_minutes < 30: + base_confidence += 0.2 + + if context.iteration_count < 5: + base_confidence += 0.1 + + return min(1.0, base_confidence) + + # Content creation methods (simplified implementations) + def _create_code_pattern_content( + self, content: str, patterns: List[str], context: ExtractionContext + ) -> str: + """Create formatted content for code pattern knowledge""" + return f"**Code Pattern: {', '.join(patterns)}**\n\n{content[:2000]}..." + + def _create_debugging_content(self, session: Dict) -> str: + """Create formatted content for debugging knowledge""" + return f"**Error:** {session.get('error_type', 'Unknown')}\n\n**Resolution:** {session.get('resolution', 'No resolution provided')}" + + def _create_process_content(self, context: ExtractionContext) -> str: + """Create formatted content for process knowledge""" + return f"**Task Type:** {context.task_data.get('task_type', 'Unknown')}\n**Iterations:** {context.iteration_count}\n**Duration:** {context.total_duration_minutes:.1f} minutes\n**Success:** {'Yes' if context.success else 'No'}" + + def _create_error_pattern_content( + self, pattern: Dict, context: ExtractionContext + ) -> str: + """Create formatted content for error pattern knowledge""" + return f"**Error Type:** {pattern['error_type']}\n**Frequency:** {pattern['frequency']}\n**Prevention:** {', '.join(pattern.get('prevention', []))}" + + def _create_optimization_content(self, context: ExtractionContext) -> str: + """Create formatted content for optimization knowledge""" + return f"**Optimization achieved in {context.total_duration_minutes:.1f} minutes with {context.iteration_count} iterations**" + + # Placeholder methods for more complex analysis functions + def _identify_debugging_sessions( + self, conversation_history: List[Dict] + ) -> List[Dict]: + """Identify debugging sessions in conversation history""" + return [] # Simplified implementation + + def _identify_error_patterns(self, conversation_history: List[Dict]) -> List[Dict]: + """Identify error patterns in conversation history""" + return [] # Simplified implementation + + def _identify_optimization_techniques( + self, context: ExtractionContext + ) -> List[str]: + """Identify optimization techniques used""" + return [] # Simplified implementation + + def _classify_problem_type(self, problem: str) -> str: + """Classify the type of problem""" + return "general" # Simplified implementation + + def _classify_solution_type(self, solution: str) -> str: + """Classify the type of solution""" + return "general" # Simplified implementation diff --git a/services/orchestrator/task_queue.py b/services/orchestrator/task_queue.py index 9178a92..87ece7d 100644 --- a/services/orchestrator/task_queue.py +++ b/services/orchestrator/task_queue.py @@ -1,124 +1,124 @@ -import asyncio -import json -import os -from typing import Any, Dict, List, Optional - -import aio_pika - -from .database import DatabaseManager - - -class TaskQueue: - def __init__(self): - self.rabbitmq_url = os.getenv( - "RABBITMQ_URL", "amqp://admin:password@rabbitmq:5672/" - ) - self.connection = None - self.channel = None - self.task_execution_engine = None # Will be set by orchestrator - - async def connect(self): - """Connect to RabbitMQ""" - if not self.connection: - self.connection = await aio_pika.connect_robust(self.rabbitmq_url) - self.channel = await self.connection.channel() - - async def assign_task(self, agent_id: str, task: dict) -> str: - """Assign a task to an agent""" - await self.connect() - - # Insert task into database - task_id = await DatabaseManager.insert_task( - title=task.get("title", "Untitled Task"), - description=task.get("description", ""), - assigned_to=agent_id, - created_by=task.get("created_by"), - ) - - # Add task_id to task data - task["id"] = task_id - task["assigned_to"] = agent_id - - # Send task to agent's queue - queue_name = f"agent_{agent_id.replace('-', '_')}" - queue = await self.channel.declare_queue(queue_name, durable=True) - - await self.channel.default_exchange.publish( - aio_pika.Message( - json.dumps(task).encode(), - delivery_mode=aio_pika.DeliveryMode.PERSISTENT, - ), - routing_key=queue_name, - ) - - return task_id - - async def list_tasks(self) -> List[Dict]: - """List all tasks""" - return await DatabaseManager.get_tasks() - - async def get_task(self, task_id: str) -> Dict: - """Get specific task""" - tasks = await self.list_tasks() - for task in tasks: - if str(task["id"]) == task_id: - return task - return None - - async def update_task_status(self, task_id: str, status: str, result: dict = None): - """Update task status""" - await DatabaseManager.update_task_status(task_id, status, result) - - async def get_pending_tasks(self) -> List[Dict]: - """Get all pending tasks""" - tasks = await self.list_tasks() - return [task for task in tasks if task["status"] == "pending"] - - async def get_agent_tasks(self, agent_id: str) -> List[Dict]: - """Get tasks assigned to specific agent""" - tasks = await self.list_tasks() - return [task for task in tasks if str(task["assigned_to"]) == agent_id] - - async def start_autonomous_execution(self, task_id: str) -> Dict[str, Any]: - """Start autonomous execution of a task""" - if not self.task_execution_engine: - raise RuntimeError("TaskExecutionEngine not configured") - - return await self.task_execution_engine.start_task_execution(task_id) - - async def get_execution_status(self, task_id: str) -> Dict[str, Any]: - """Get execution status of a task""" - if not self.task_execution_engine: - raise RuntimeError("TaskExecutionEngine not configured") - - return await self.task_execution_engine.get_execution_status(task_id) - - async def get_task_iterations(self, task_id: str) -> List[Dict[str, Any]]: - """Get task iteration history""" - if not self.task_execution_engine: - raise RuntimeError("TaskExecutionEngine not configured") - - return await self.task_execution_engine.get_task_iterations(task_id) - - async def handle_human_response(self, task_id: str, response: str) -> bool: - """Handle human response to a task question""" - if not self.task_execution_engine: - raise RuntimeError("TaskExecutionEngine not configured") - - return await self.task_execution_engine.handle_human_response(task_id, response) - - async def cancel_task_execution(self, task_id: str) -> bool: - """Cancel autonomous execution of a task""" - if not self.task_execution_engine: - raise RuntimeError("TaskExecutionEngine not configured") - - return await self.task_execution_engine.cancel_task_execution(task_id) - - def set_task_execution_engine(self, engine): - """Set the task execution engine reference""" - self.task_execution_engine = engine - - async def close(self): - """Close RabbitMQ connection""" - if self.connection: - await self.connection.close() +import asyncio +import json +import os +from typing import Any, Dict, List, Optional + +import aio_pika + +from .database import DatabaseManager + + +class TaskQueue: + def __init__(self): + self.rabbitmq_url = os.getenv( + "RABBITMQ_URL", "amqp://admin:password@rabbitmq:5672/" + ) + self.connection = None + self.channel = None + self.task_execution_engine = None # Will be set by orchestrator + + async def connect(self): + """Connect to RabbitMQ""" + if not self.connection: + self.connection = await aio_pika.connect_robust(self.rabbitmq_url) + self.channel = await self.connection.channel() + + async def assign_task(self, agent_id: str, task: dict) -> str: + """Assign a task to an agent""" + await self.connect() + + # Insert task into database + task_id = await DatabaseManager.insert_task( + title=task.get("title", "Untitled Task"), + description=task.get("description", ""), + assigned_to=agent_id, + created_by=task.get("created_by"), + ) + + # Add task_id to task data + task["id"] = task_id + task["assigned_to"] = agent_id + + # Send task to agent's queue + queue_name = f"agent_{agent_id.replace('-', '_')}" + queue = await self.channel.declare_queue(queue_name, durable=True) + + await self.channel.default_exchange.publish( + aio_pika.Message( + json.dumps(task).encode(), + delivery_mode=aio_pika.DeliveryMode.PERSISTENT, + ), + routing_key=queue_name, + ) + + return task_id + + async def list_tasks(self) -> List[Dict]: + """List all tasks""" + return await DatabaseManager.get_tasks() + + async def get_task(self, task_id: str) -> Dict: + """Get specific task""" + tasks = await self.list_tasks() + for task in tasks: + if str(task["id"]) == task_id: + return task + return None + + async def update_task_status(self, task_id: str, status: str, result: dict = None): + """Update task status""" + await DatabaseManager.update_task_status(task_id, status, result) + + async def get_pending_tasks(self) -> List[Dict]: + """Get all pending tasks""" + tasks = await self.list_tasks() + return [task for task in tasks if task["status"] == "pending"] + + async def get_agent_tasks(self, agent_id: str) -> List[Dict]: + """Get tasks assigned to specific agent""" + tasks = await self.list_tasks() + return [task for task in tasks if str(task["assigned_to"]) == agent_id] + + async def start_autonomous_execution(self, task_id: str) -> Dict[str, Any]: + """Start autonomous execution of a task""" + if not self.task_execution_engine: + raise RuntimeError("TaskExecutionEngine not configured") + + return await self.task_execution_engine.start_task_execution(task_id) + + async def get_execution_status(self, task_id: str) -> Dict[str, Any]: + """Get execution status of a task""" + if not self.task_execution_engine: + raise RuntimeError("TaskExecutionEngine not configured") + + return await self.task_execution_engine.get_execution_status(task_id) + + async def get_task_iterations(self, task_id: str) -> List[Dict[str, Any]]: + """Get task iteration history""" + if not self.task_execution_engine: + raise RuntimeError("TaskExecutionEngine not configured") + + return await self.task_execution_engine.get_task_iterations(task_id) + + async def handle_human_response(self, task_id: str, response: str) -> bool: + """Handle human response to a task question""" + if not self.task_execution_engine: + raise RuntimeError("TaskExecutionEngine not configured") + + return await self.task_execution_engine.handle_human_response(task_id, response) + + async def cancel_task_execution(self, task_id: str) -> bool: + """Cancel autonomous execution of a task""" + if not self.task_execution_engine: + raise RuntimeError("TaskExecutionEngine not configured") + + return await self.task_execution_engine.cancel_task_execution(task_id) + + def set_task_execution_engine(self, engine): + """Set the task execution engine reference""" + self.task_execution_engine = engine + + async def close(self): + """Close RabbitMQ connection""" + if self.connection: + await self.connection.close() diff --git a/services/orchestrator/team_knowledge_manager.py b/services/orchestrator/team_knowledge_manager.py index a20741b..f7db4ab 100644 --- a/services/orchestrator/team_knowledge_manager.py +++ b/services/orchestrator/team_knowledge_manager.py @@ -1,869 +1,869 @@ -""" -Team Knowledge Manager for FuzeAgent - -This module manages team-level knowledge aggregation, filtering organization knowledge -for team relevance, and facilitating knowledge sharing between agents within teams. -""" - -import asyncio -import json -import logging -import uuid -from dataclasses import dataclass -from datetime import datetime, timedelta -from typing import Any, Dict, List, Optional, Tuple - -import asyncpg -from sentence_transformers import SentenceTransformer - -from .organization_rag_manager import ( - ContentType, - KnowledgeCategory, - KnowledgeSearchResult, - OrganizationRAGManager, - SourceType, - VisibilityLevel, -) - -logger = logging.getLogger(__name__) - - -@dataclass -class TeamKnowledge: - """Represents team-level knowledge""" - - id: str - team_id: str - organization_id: str - title: str - content: str - content_type: ContentType - knowledge_category: KnowledgeCategory - embedding: Optional[List[float]] - source_type: SourceType - contributing_agents: List[str] - source_knowledge_ids: List[str] - aggregation_method: str - team_relevance_score: float - agent_adoption_rate: float - effectiveness_score: float - visibility_level: VisibilityLevel - metadata: Dict[str, Any] - tags: List[str] - created_at: datetime - updated_at: datetime - last_accessed: Optional[datetime] - - -@dataclass -class TeamKnowledgeSearchResult: - """Result of team knowledge search""" - - team_knowledge: TeamKnowledge - similarity_score: float - relevance_score: float - team_fit_score: float - combined_score: float - - -class TeamKnowledgeManager: - """ - Manages team-specific knowledge base with intelligent aggregation - from organization knowledge and agent contributions. - """ - - def __init__( - self, database_url: str, organization_rag_manager: OrganizationRAGManager - ): - self.database_url = database_url - self.org_rag_manager = organization_rag_manager - self.pool: Optional[asyncpg.Pool] = None - - # Initialize embedding model - self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") - self.embedding_dim = 384 - - # Configuration - self.min_team_relevance = 0.4 - self.adoption_threshold = 0.6 # 60% of team agents should find it useful - self.effectiveness_decay_days = 30 - - # Statistics - self.team_queries_processed = 0 - self.team_knowledge_created = 0 - self.aggregations_performed = 0 - - async def initialize(self): - """Initialize the team knowledge manager""" - logger.info("Initializing TeamKnowledgeManager") - - try: - self.pool = await asyncpg.create_pool( - self.database_url, min_size=2, max_size=10, command_timeout=60 - ) - - logger.info("TeamKnowledgeManager initialized successfully") - - except Exception as e: - logger.error(f"Failed to initialize TeamKnowledgeManager: {e}") - raise - - async def close(self): - """Close database connections""" - if self.pool: - await self.pool.close() - logger.info("TeamKnowledgeManager closed") - - async def create_team_knowledge( - self, - team_id: str, - title: str, - content: str, - content_type: ContentType = ContentType.TEXT, - knowledge_category: KnowledgeCategory = KnowledgeCategory.DEVELOPMENT, - source_type: SourceType = SourceType.TEAM_AGGREGATION, - contributing_agents: Optional[List[str]] = None, - source_knowledge_ids: Optional[List[str]] = None, - aggregation_method: str = "synthesis", - team_relevance_score: float = 0.7, - metadata: Optional[Dict[str, Any]] = None, - tags: Optional[List[str]] = None, - ) -> str: - """Create team-specific knowledge""" - - team_knowledge_id = str(uuid.uuid4()) - embedding = self._generate_embedding(content) - - async with self.pool.acquire() as conn: - # Get organization_id for this team - org_id = await conn.fetchval( - """ - SELECT organization_id FROM teams WHERE id = $1 - """, - team_id, - ) - - if not org_id: - raise ValueError(f"Team {team_id} not found") - - await conn.execute( - """ - INSERT INTO team_knowledge_base ( - id, team_id, organization_id, title, content, content_type, - knowledge_category, embedding, source_type, contributing_agents, - source_knowledge_ids, aggregation_method, team_relevance_score, - metadata, tags - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15) - """, - team_knowledge_id, - team_id, - org_id, - title, - content, - content_type.value, - knowledge_category.value, - embedding, - source_type.value, - contributing_agents or [], - source_knowledge_ids or [], - aggregation_method, - team_relevance_score, - json.dumps(metadata or {}), - tags or [], - ) - - self.team_knowledge_created += 1 - - logger.info(f"Created team knowledge {team_knowledge_id} for team {team_id}") - return team_knowledge_id - - async def search_team_knowledge( - self, - team_id: str, - query: str, - categories: Optional[List[KnowledgeCategory]] = None, - content_types: Optional[List[ContentType]] = None, - include_org_knowledge: bool = True, - limit: int = 10, - min_similarity: float = 0.3, - ) -> List[TeamKnowledgeSearchResult]: - """Search team knowledge with optional organization knowledge inclusion""" - - self.team_queries_processed += 1 - query_embedding = self._generate_embedding(query) - results = [] - - async with self.pool.acquire() as conn: - # Search team-specific knowledge - team_results = await self._search_team_specific_knowledge( - conn, - team_id, - query_embedding, - categories, - content_types, - limit, - min_similarity, - ) - results.extend(team_results) - - # Search organization knowledge filtered for team relevance - if include_org_knowledge and len(results) < limit: - org_results = await self._search_org_knowledge_for_team( - conn, - team_id, - query, - categories, - content_types, - limit - len(results), - min_similarity, - ) - results.extend(org_results) - - # Sort by combined score - results.sort(key=lambda x: x.combined_score, reverse=True) - return results[:limit] - - async def aggregate_agent_knowledge_to_team( - self, - team_id: str, - agent_id: str, - agent_memory_ids: List[str], - aggregation_method: str = "synthesis", - ) -> Optional[str]: - """Aggregate multiple agent memories into team knowledge""" - - async with self.pool.acquire() as conn: - # Get agent memories - agent_memories = await conn.fetch( - """ - SELECT * FROM agent_memory - WHERE id = ANY($1) AND agent_id = $2 - ORDER BY confidence_score DESC, created_at DESC - """, - agent_memory_ids, - agent_id, - ) - - if not agent_memories: - return None - - # Analyze memories for commonalities - analysis_result = await self._analyze_memories_for_aggregation( - agent_memories - ) - - if analysis_result["aggregation_value"] < self.min_team_relevance: - logger.debug( - f"Agent memories don't meet team relevance threshold: {analysis_result['aggregation_value']}" - ) - return None - - # Create aggregated knowledge - team_knowledge_id = await self.create_team_knowledge( - team_id=team_id, - title=analysis_result["title"], - content=analysis_result["content"], - content_type=analysis_result["content_type"], - knowledge_category=analysis_result["category"], - source_type=SourceType.AGENT_CONTRIBUTION, - contributing_agents=[agent_id], - aggregation_method=aggregation_method, - team_relevance_score=analysis_result["aggregation_value"], - metadata=analysis_result["metadata"], - tags=analysis_result["tags"], - ) - - # Mark original memories as aggregated - await conn.execute( - """ - UPDATE agent_memory - SET propagated_to_team = TRUE, team_context_id = $2 - WHERE id = ANY($1) - """, - agent_memory_ids, - team_knowledge_id, - ) - - self.aggregations_performed += 1 - - logger.info( - f"Aggregated {len(agent_memory_ids)} agent memories into team knowledge {team_knowledge_id}" - ) - return team_knowledge_id - - async def get_team_knowledge_context( - self, - team_id: str, - task_context: Dict[str, Any], - agent_id: Optional[str] = None, - max_context_items: int = 5, - ) -> Dict[str, Any]: - """Get relevant team knowledge for task execution context""" - - # Build context query from task information - context_query = self._build_context_query(task_context) - - # Search for relevant knowledge - search_results = await self.search_team_knowledge( - team_id=team_id, - query=context_query, - limit=max_context_items, - min_similarity=0.4, - ) - - # Get team statistics - team_stats = await self.get_team_knowledge_stats(team_id) - - # Build context - context = { - "team_id": team_id, - "relevant_knowledge": [ - { - "id": result.team_knowledge.id, - "title": result.team_knowledge.title, - "content": ( - result.team_knowledge.content[:500] + "..." - if len(result.team_knowledge.content) > 500 - else result.team_knowledge.content - ), - "category": result.team_knowledge.knowledge_category.value, - "relevance_score": result.combined_score, - "usage_stats": { - "adoption_rate": result.team_knowledge.agent_adoption_rate, - "effectiveness": result.team_knowledge.effectiveness_score, - }, - } - for result in search_results - ], - "team_knowledge_stats": team_stats, - "context_query": context_query, - "generated_at": datetime.now().isoformat(), - } - - return context - - async def update_knowledge_effectiveness( - self, - team_knowledge_id: str, - agent_id: str, - task_success: bool, - feedback_score: Optional[float] = None, - usage_context: Optional[Dict[str, Any]] = None, - ): - """Update knowledge effectiveness based on agent usage""" - - async with self.pool.acquire() as conn: - # Get current knowledge - knowledge = await conn.fetchrow( - """ - SELECT * FROM team_knowledge_base WHERE id = $1 - """, - team_knowledge_id, - ) - - if not knowledge: - return - - # Calculate new effectiveness score - success_weight = 1.0 if task_success else -0.3 - feedback_weight = (feedback_score or 0.5) - 0.5 - - # Update effectiveness with exponential moving average - current_effectiveness = knowledge["effectiveness_score"] - new_effectiveness = ( - current_effectiveness * 0.8 + (success_weight + feedback_weight) * 0.2 - ) - new_effectiveness = max(0.0, min(1.0, new_effectiveness)) - - # Update agent adoption tracking - contributing_agents = knowledge["contributing_agents"] or [] - if agent_id not in contributing_agents: - contributing_agents.append(agent_id) - - # Calculate adoption rate (agents who used it / total team agents) - team_agent_count = await conn.fetchval( - """ - SELECT COUNT(*) FROM agents WHERE team_id = $1 - """, - knowledge["team_id"], - ) - - adoption_rate = len(contributing_agents) / max(1, team_agent_count) - - # Update knowledge - await conn.execute( - """ - UPDATE team_knowledge_base - SET effectiveness_score = $2, - agent_adoption_rate = $3, - contributing_agents = $4, - last_accessed = NOW(), - updated_at = NOW() - WHERE id = $1 - """, - team_knowledge_id, - new_effectiveness, - adoption_rate, - contributing_agents, - ) - - logger.debug( - f"Updated knowledge {team_knowledge_id} effectiveness: {new_effectiveness:.2f}, adoption: {adoption_rate:.2f}" - ) - - async def get_team_knowledge_stats(self, team_id: str) -> Dict[str, Any]: - """Get comprehensive team knowledge statistics""" - - async with self.pool.acquire() as conn: - # Basic statistics - basic_stats = await conn.fetchrow( - """ - SELECT - COUNT(*) as total_knowledge, - COUNT(DISTINCT knowledge_category) as categories, - COUNT(DISTINCT unnest(contributing_agents)) as contributing_agents, - AVG(team_relevance_score) as avg_relevance, - AVG(effectiveness_score) as avg_effectiveness, - AVG(agent_adoption_rate) as avg_adoption_rate - FROM team_knowledge_base - WHERE team_id = $1 - """, - team_id, - ) - - # Category breakdown - category_stats = await conn.fetch( - """ - SELECT - knowledge_category, - COUNT(*) as count, - AVG(effectiveness_score) as avg_effectiveness, - AVG(agent_adoption_rate) as avg_adoption - FROM team_knowledge_base - WHERE team_id = $1 - GROUP BY knowledge_category - ORDER BY count DESC - """, - team_id, - ) - - # Most effective knowledge - top_knowledge = await conn.fetch( - """ - SELECT - title, - knowledge_category, - effectiveness_score, - agent_adoption_rate - FROM team_knowledge_base - WHERE team_id = $1 - ORDER BY effectiveness_score DESC - LIMIT 5 - """, - team_id, - ) - - return { - "team_id": team_id, - "basic_stats": dict(basic_stats) if basic_stats else {}, - "category_breakdown": [dict(cat) for cat in category_stats], - "top_knowledge": [dict(know) for know in top_knowledge], - "generated_at": datetime.now().isoformat(), - } - - async def _search_team_specific_knowledge( - self, - conn, - team_id: str, - query_embedding: List[float], - categories: Optional[List[KnowledgeCategory]], - content_types: Optional[List[ContentType]], - limit: int, - min_similarity: float, - ) -> List[TeamKnowledgeSearchResult]: - """Search team-specific knowledge base""" - - # Build query conditions - where_conditions = ["team_id = $2"] - params = [query_embedding, team_id] - param_idx = 3 - - if categories: - where_conditions.append(f"knowledge_category = ANY(${param_idx})") - params.append([cat.value for cat in categories]) - param_idx += 1 - - if content_types: - where_conditions.append(f"content_type = ANY(${param_idx})") - params.append([ct.value for ct in content_types]) - param_idx += 1 - - where_conditions.append(f"(1 - (embedding <=> $1)) >= ${param_idx}") - params.append(min_similarity) - param_idx += 1 - - where_clause = "WHERE " + " AND ".join(where_conditions) - - results = await conn.fetch( - f""" - SELECT - *, - (1 - (embedding <=> $1)) as similarity_score - FROM team_knowledge_base - {where_clause} - ORDER BY similarity_score DESC, effectiveness_score DESC - LIMIT ${param_idx} - """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params - *params, - limit, - ) - - search_results = [] - for row in results: - team_knowledge = self._row_to_team_knowledge(row) - - # Calculate team fit score based on adoption and effectiveness - team_fit_score = ( - team_knowledge.agent_adoption_rate * 0.4 - + team_knowledge.effectiveness_score * 0.6 - ) - - combined_score = ( - float(row["similarity_score"]) * 0.4 - + team_knowledge.team_relevance_score * 0.3 - + team_fit_score * 0.3 - ) - - search_results.append( - TeamKnowledgeSearchResult( - team_knowledge=team_knowledge, - similarity_score=float(row["similarity_score"]), - relevance_score=team_knowledge.team_relevance_score, - team_fit_score=team_fit_score, - combined_score=combined_score, - ) - ) - - return search_results - - async def _search_org_knowledge_for_team( - self, - conn, - team_id: str, - query: str, - categories: Optional[List[KnowledgeCategory]], - content_types: Optional[List[ContentType]], - limit: int, - min_similarity: float, - ) -> List[TeamKnowledgeSearchResult]: - """Search organization knowledge filtered for team relevance""" - - # Get organization ID for the team - org_id = await conn.fetchval( - """ - SELECT organization_id FROM teams WHERE id = $1 - """, - team_id, - ) - - if not org_id: - return [] - - # Search organization knowledge - org_results = await self.org_rag_manager.search_knowledge( - organization_id=str(org_id), - query=query, - categories=categories, - content_types=content_types, - limit=limit * 2, # Get more to filter for team relevance - min_similarity=min_similarity, - requester_team_id=team_id, - ) - - # Convert to team knowledge search results with team relevance scoring - team_results = [] - for org_result in org_results: - # Calculate team relevance based on source and usage - team_relevance = await self._calculate_team_relevance( - conn, team_id, org_result.knowledge - ) - - if team_relevance >= self.min_team_relevance: - # Create pseudo team knowledge for consistent interface - pseudo_team_knowledge = TeamKnowledge( - id=org_result.knowledge.id, - team_id=team_id, - organization_id=org_result.knowledge.organization_id, - title=org_result.knowledge.title, - content=org_result.knowledge.content, - content_type=org_result.knowledge.content_type, - knowledge_category=org_result.knowledge.knowledge_category, - embedding=org_result.knowledge.embedding, - source_type=org_result.knowledge.source_type, - contributing_agents=[], - source_knowledge_ids=[org_result.knowledge.id], - aggregation_method="organization_filter", - team_relevance_score=team_relevance, - agent_adoption_rate=0.0, - effectiveness_score=org_result.knowledge.success_correlation, - visibility_level=org_result.knowledge.visibility_level, - metadata=org_result.knowledge.metadata, - tags=org_result.knowledge.tags, - created_at=org_result.knowledge.created_at, - updated_at=org_result.knowledge.updated_at, - last_accessed=org_result.knowledge.last_accessed, - ) - - combined_score = ( - org_result.similarity_score * 0.5 + team_relevance * 0.5 - ) - - team_results.append( - TeamKnowledgeSearchResult( - team_knowledge=pseudo_team_knowledge, - similarity_score=org_result.similarity_score, - relevance_score=org_result.relevance_score, - team_fit_score=team_relevance, - combined_score=combined_score, - ) - ) - - return team_results[:limit] - - async def _calculate_team_relevance( - self, conn, team_id: str, org_knowledge - ) -> float: - """Calculate how relevant organization knowledge is for a specific team""" - - relevance_factors = [] - - # Factor 1: Source team match - if org_knowledge.source_team_id == team_id: - relevance_factors.append(1.0) - elif org_knowledge.source_team_id: - # Check if source team is similar to current team - team_similarity = await self._calculate_team_similarity( - conn, team_id, org_knowledge.source_team_id - ) - relevance_factors.append(team_similarity) - else: - relevance_factors.append(0.3) # No team context - - # Factor 2: Category relevance to team's work - team_categories = await self._get_team_primary_categories(conn, team_id) - if org_knowledge.knowledge_category.value in team_categories: - relevance_factors.append(0.9) - else: - relevance_factors.append(0.4) - - # Factor 3: Usage by team agents - team_usage = ( - await conn.fetchval( - """ - SELECT COUNT(DISTINCT source_agent_id)::float / NULLIF( - (SELECT COUNT(*) FROM agents WHERE team_id = $1), 0 - ) - FROM organization_knowledge_base - WHERE id = $2 AND source_agent_id IN ( - SELECT id FROM agents WHERE team_id = $1 - ) - """, - team_id, - org_knowledge.id, - ) - or 0.0 - ) - relevance_factors.append(team_usage) - - # Factor 4: Base quality and relevance - relevance_factors.append(org_knowledge.quality_score) - relevance_factors.append(org_knowledge.relevance_score) - - # Calculate weighted average - weights = [0.3, 0.25, 0.25, 0.1, 0.1] - team_relevance = sum(f * w for f, w in zip(relevance_factors, weights)) - - return min(1.0, max(0.0, team_relevance)) - - async def _calculate_team_similarity( - self, conn, team_id1: str, team_id2: str - ) -> float: - """Calculate similarity between two teams based on their work patterns""" - - # Simple implementation based on team type and settings - team_info = await conn.fetch( - """ - SELECT id, team_type, settings FROM teams - WHERE id IN ($1, $2) - """, - team_id1, - team_id2, - ) - - if len(team_info) != 2: - return 0.0 - - team1, team2 = team_info - - # Type similarity - type_similarity = 1.0 if team1["team_type"] == team2["team_type"] else 0.5 - - # Settings similarity (simplified) - settings1 = team1["settings"] or {} - settings2 = team2["settings"] or {} - - common_keys = set(settings1.keys()) & set(settings2.keys()) - if common_keys: - settings_similarity = sum( - 1.0 if settings1.get(key) == settings2.get(key) else 0.0 - for key in common_keys - ) / len(common_keys) - else: - settings_similarity = 0.5 - - return type_similarity * 0.7 + settings_similarity * 0.3 - - async def _get_team_primary_categories(self, conn, team_id: str) -> List[str]: - """Get primary knowledge categories this team works with""" - - categories = await conn.fetch( - """ - SELECT knowledge_category, COUNT(*) as usage_count - FROM team_knowledge_base - WHERE team_id = $1 - GROUP BY knowledge_category - ORDER BY usage_count DESC - LIMIT 3 - """, - team_id, - ) - - return [cat["knowledge_category"] for cat in categories] - - def _generate_embedding(self, text: str) -> List[float]: - """Generate embedding for text using sentence transformers""" - try: - embedding = self.embedding_model.encode(text, convert_to_tensor=False) - return embedding.tolist() - except Exception as e: - logger.error(f"Error generating embedding: {e}") - return [0.0] * self.embedding_dim - - def _row_to_team_knowledge(self, row) -> TeamKnowledge: - """Convert database row to TeamKnowledge object""" - return TeamKnowledge( - id=str(row["id"]), - team_id=str(row["team_id"]), - organization_id=str(row["organization_id"]), - title=row["title"], - content=row["content"], - content_type=ContentType(row["content_type"]), - knowledge_category=KnowledgeCategory(row["knowledge_category"]), - embedding=row["embedding"] if row["embedding"] else None, - source_type=SourceType(row["source_type"]), - contributing_agents=row["contributing_agents"] or [], - source_knowledge_ids=row["source_knowledge_ids"] or [], - aggregation_method=row["aggregation_method"], - team_relevance_score=row["team_relevance_score"], - agent_adoption_rate=row["agent_adoption_rate"], - effectiveness_score=row["effectiveness_score"], - visibility_level=VisibilityLevel(row["visibility_level"]), - metadata=( - json.loads(row["metadata"]) - if isinstance(row["metadata"], str) - else row["metadata"] - ), - tags=row["tags"] or [], - created_at=row["created_at"], - updated_at=row["updated_at"], - last_accessed=row["last_accessed"], - ) - - def _build_context_query(self, task_context: Dict[str, Any]) -> str: - """Build a search query from task context""" - query_parts = [] - - if task_context.get("task_type"): - query_parts.append(task_context["task_type"]) - - if task_context.get("description"): - query_parts.append(task_context["description"]) - - if task_context.get("technologies"): - query_parts.extend(task_context["technologies"]) - - if task_context.get("domain"): - query_parts.append(task_context["domain"]) - - return " ".join(query_parts) - - async def _analyze_memories_for_aggregation( - self, agent_memories: List - ) -> Dict[str, Any]: - """Analyze agent memories to determine if they should be aggregated""" - - if not agent_memories: - return {"aggregation_value": 0.0} - - # Simple aggregation analysis - # In practice, this could use more sophisticated NLP - - # Calculate average confidence and success correlation - avg_confidence = sum(mem["confidence_score"] for mem in agent_memories) / len( - agent_memories - ) - avg_success = sum( - mem.get("success_correlation", 0.0) for mem in agent_memories - ) / len(agent_memories) - - # Find common themes - all_content = " ".join(mem["content"] for mem in agent_memories) - - # Determine primary category - categories = [mem.get("memory_type", "general") for mem in agent_memories] - primary_category = ( - max(set(categories), key=categories.count) if categories else "general" - ) - - # Create aggregated content (simplified) - title = f"Team Knowledge: {primary_category.replace('_', ' ').title()}" - content = ( - f"Aggregated knowledge from {len(agent_memories)} agent experiences:\n\n" - + all_content[:1000] - ) - - # Map memory type to knowledge category - category_mapping = { - "code_pattern": KnowledgeCategory.DEVELOPMENT, - "task_outcome": KnowledgeCategory.PROCESS, - "debugging": KnowledgeCategory.TROUBLESHOOTING, - "optimization": KnowledgeCategory.DEVELOPMENT, - "testing": KnowledgeCategory.TESTING, - } - - knowledge_category = category_mapping.get( - primary_category, KnowledgeCategory.DEVELOPMENT - ) - - # Determine content type - content_type = ( - ContentType.CODE if "code" in primary_category else ContentType.TEXT - ) - - # Calculate aggregation value - aggregation_value = min(1.0, (avg_confidence + avg_success) / 2.0) - - return { - "aggregation_value": aggregation_value, - "title": title, - "content": content, - "content_type": content_type, - "category": knowledge_category, - "metadata": { - "source_memory_count": len(agent_memories), - "avg_confidence": avg_confidence, - "avg_success_correlation": avg_success, - "primary_type": primary_category, - }, - "tags": [primary_category, "aggregated", "agent_contribution"], - } +""" +Team Knowledge Manager for FuzeAgent + +This module manages team-level knowledge aggregation, filtering organization knowledge +for team relevance, and facilitating knowledge sharing between agents within teams. +""" + +import asyncio +import json +import logging +import uuid +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Any, Dict, List, Optional, Tuple + +import asyncpg +from sentence_transformers import SentenceTransformer + +from .organization_rag_manager import ( + ContentType, + KnowledgeCategory, + KnowledgeSearchResult, + OrganizationRAGManager, + SourceType, + VisibilityLevel, +) + +logger = logging.getLogger(__name__) + + +@dataclass +class TeamKnowledge: + """Represents team-level knowledge""" + + id: str + team_id: str + organization_id: str + title: str + content: str + content_type: ContentType + knowledge_category: KnowledgeCategory + embedding: Optional[List[float]] + source_type: SourceType + contributing_agents: List[str] + source_knowledge_ids: List[str] + aggregation_method: str + team_relevance_score: float + agent_adoption_rate: float + effectiveness_score: float + visibility_level: VisibilityLevel + metadata: Dict[str, Any] + tags: List[str] + created_at: datetime + updated_at: datetime + last_accessed: Optional[datetime] + + +@dataclass +class TeamKnowledgeSearchResult: + """Result of team knowledge search""" + + team_knowledge: TeamKnowledge + similarity_score: float + relevance_score: float + team_fit_score: float + combined_score: float + + +class TeamKnowledgeManager: + """ + Manages team-specific knowledge base with intelligent aggregation + from organization knowledge and agent contributions. + """ + + def __init__( + self, database_url: str, organization_rag_manager: OrganizationRAGManager + ): + self.database_url = database_url + self.org_rag_manager = organization_rag_manager + self.pool: Optional[asyncpg.Pool] = None + + # Initialize embedding model + self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") + self.embedding_dim = 384 + + # Configuration + self.min_team_relevance = 0.4 + self.adoption_threshold = 0.6 # 60% of team agents should find it useful + self.effectiveness_decay_days = 30 + + # Statistics + self.team_queries_processed = 0 + self.team_knowledge_created = 0 + self.aggregations_performed = 0 + + async def initialize(self): + """Initialize the team knowledge manager""" + logger.info("Initializing TeamKnowledgeManager") + + try: + self.pool = await asyncpg.create_pool( + self.database_url, min_size=2, max_size=10, command_timeout=60 + ) + + logger.info("TeamKnowledgeManager initialized successfully") + + except Exception as e: + logger.error(f"Failed to initialize TeamKnowledgeManager: {e}") + raise + + async def close(self): + """Close database connections""" + if self.pool: + await self.pool.close() + logger.info("TeamKnowledgeManager closed") + + async def create_team_knowledge( + self, + team_id: str, + title: str, + content: str, + content_type: ContentType = ContentType.TEXT, + knowledge_category: KnowledgeCategory = KnowledgeCategory.DEVELOPMENT, + source_type: SourceType = SourceType.TEAM_AGGREGATION, + contributing_agents: Optional[List[str]] = None, + source_knowledge_ids: Optional[List[str]] = None, + aggregation_method: str = "synthesis", + team_relevance_score: float = 0.7, + metadata: Optional[Dict[str, Any]] = None, + tags: Optional[List[str]] = None, + ) -> str: + """Create team-specific knowledge""" + + team_knowledge_id = str(uuid.uuid4()) + embedding = self._generate_embedding(content) + + async with self.pool.acquire() as conn: + # Get organization_id for this team + org_id = await conn.fetchval( + """ + SELECT organization_id FROM teams WHERE id = $1 + """, + team_id, + ) + + if not org_id: + raise ValueError(f"Team {team_id} not found") + + await conn.execute( + """ + INSERT INTO team_knowledge_base ( + id, team_id, organization_id, title, content, content_type, + knowledge_category, embedding, source_type, contributing_agents, + source_knowledge_ids, aggregation_method, team_relevance_score, + metadata, tags + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15) + """, + team_knowledge_id, + team_id, + org_id, + title, + content, + content_type.value, + knowledge_category.value, + embedding, + source_type.value, + contributing_agents or [], + source_knowledge_ids or [], + aggregation_method, + team_relevance_score, + json.dumps(metadata or {}), + tags or [], + ) + + self.team_knowledge_created += 1 + + logger.info(f"Created team knowledge {team_knowledge_id} for team {team_id}") + return team_knowledge_id + + async def search_team_knowledge( + self, + team_id: str, + query: str, + categories: Optional[List[KnowledgeCategory]] = None, + content_types: Optional[List[ContentType]] = None, + include_org_knowledge: bool = True, + limit: int = 10, + min_similarity: float = 0.3, + ) -> List[TeamKnowledgeSearchResult]: + """Search team knowledge with optional organization knowledge inclusion""" + + self.team_queries_processed += 1 + query_embedding = self._generate_embedding(query) + results = [] + + async with self.pool.acquire() as conn: + # Search team-specific knowledge + team_results = await self._search_team_specific_knowledge( + conn, + team_id, + query_embedding, + categories, + content_types, + limit, + min_similarity, + ) + results.extend(team_results) + + # Search organization knowledge filtered for team relevance + if include_org_knowledge and len(results) < limit: + org_results = await self._search_org_knowledge_for_team( + conn, + team_id, + query, + categories, + content_types, + limit - len(results), + min_similarity, + ) + results.extend(org_results) + + # Sort by combined score + results.sort(key=lambda x: x.combined_score, reverse=True) + return results[:limit] + + async def aggregate_agent_knowledge_to_team( + self, + team_id: str, + agent_id: str, + agent_memory_ids: List[str], + aggregation_method: str = "synthesis", + ) -> Optional[str]: + """Aggregate multiple agent memories into team knowledge""" + + async with self.pool.acquire() as conn: + # Get agent memories + agent_memories = await conn.fetch( + """ + SELECT * FROM agent_memory + WHERE id = ANY($1) AND agent_id = $2 + ORDER BY confidence_score DESC, created_at DESC + """, + agent_memory_ids, + agent_id, + ) + + if not agent_memories: + return None + + # Analyze memories for commonalities + analysis_result = await self._analyze_memories_for_aggregation( + agent_memories + ) + + if analysis_result["aggregation_value"] < self.min_team_relevance: + logger.debug( + f"Agent memories don't meet team relevance threshold: {analysis_result['aggregation_value']}" + ) + return None + + # Create aggregated knowledge + team_knowledge_id = await self.create_team_knowledge( + team_id=team_id, + title=analysis_result["title"], + content=analysis_result["content"], + content_type=analysis_result["content_type"], + knowledge_category=analysis_result["category"], + source_type=SourceType.AGENT_CONTRIBUTION, + contributing_agents=[agent_id], + aggregation_method=aggregation_method, + team_relevance_score=analysis_result["aggregation_value"], + metadata=analysis_result["metadata"], + tags=analysis_result["tags"], + ) + + # Mark original memories as aggregated + await conn.execute( + """ + UPDATE agent_memory + SET propagated_to_team = TRUE, team_context_id = $2 + WHERE id = ANY($1) + """, + agent_memory_ids, + team_knowledge_id, + ) + + self.aggregations_performed += 1 + + logger.info( + f"Aggregated {len(agent_memory_ids)} agent memories into team knowledge {team_knowledge_id}" + ) + return team_knowledge_id + + async def get_team_knowledge_context( + self, + team_id: str, + task_context: Dict[str, Any], + agent_id: Optional[str] = None, + max_context_items: int = 5, + ) -> Dict[str, Any]: + """Get relevant team knowledge for task execution context""" + + # Build context query from task information + context_query = self._build_context_query(task_context) + + # Search for relevant knowledge + search_results = await self.search_team_knowledge( + team_id=team_id, + query=context_query, + limit=max_context_items, + min_similarity=0.4, + ) + + # Get team statistics + team_stats = await self.get_team_knowledge_stats(team_id) + + # Build context + context = { + "team_id": team_id, + "relevant_knowledge": [ + { + "id": result.team_knowledge.id, + "title": result.team_knowledge.title, + "content": ( + result.team_knowledge.content[:500] + "..." + if len(result.team_knowledge.content) > 500 + else result.team_knowledge.content + ), + "category": result.team_knowledge.knowledge_category.value, + "relevance_score": result.combined_score, + "usage_stats": { + "adoption_rate": result.team_knowledge.agent_adoption_rate, + "effectiveness": result.team_knowledge.effectiveness_score, + }, + } + for result in search_results + ], + "team_knowledge_stats": team_stats, + "context_query": context_query, + "generated_at": datetime.now().isoformat(), + } + + return context + + async def update_knowledge_effectiveness( + self, + team_knowledge_id: str, + agent_id: str, + task_success: bool, + feedback_score: Optional[float] = None, + usage_context: Optional[Dict[str, Any]] = None, + ): + """Update knowledge effectiveness based on agent usage""" + + async with self.pool.acquire() as conn: + # Get current knowledge + knowledge = await conn.fetchrow( + """ + SELECT * FROM team_knowledge_base WHERE id = $1 + """, + team_knowledge_id, + ) + + if not knowledge: + return + + # Calculate new effectiveness score + success_weight = 1.0 if task_success else -0.3 + feedback_weight = (feedback_score or 0.5) - 0.5 + + # Update effectiveness with exponential moving average + current_effectiveness = knowledge["effectiveness_score"] + new_effectiveness = ( + current_effectiveness * 0.8 + (success_weight + feedback_weight) * 0.2 + ) + new_effectiveness = max(0.0, min(1.0, new_effectiveness)) + + # Update agent adoption tracking + contributing_agents = knowledge["contributing_agents"] or [] + if agent_id not in contributing_agents: + contributing_agents.append(agent_id) + + # Calculate adoption rate (agents who used it / total team agents) + team_agent_count = await conn.fetchval( + """ + SELECT COUNT(*) FROM agents WHERE team_id = $1 + """, + knowledge["team_id"], + ) + + adoption_rate = len(contributing_agents) / max(1, team_agent_count) + + # Update knowledge + await conn.execute( + """ + UPDATE team_knowledge_base + SET effectiveness_score = $2, + agent_adoption_rate = $3, + contributing_agents = $4, + last_accessed = NOW(), + updated_at = NOW() + WHERE id = $1 + """, + team_knowledge_id, + new_effectiveness, + adoption_rate, + contributing_agents, + ) + + logger.debug( + f"Updated knowledge {team_knowledge_id} effectiveness: {new_effectiveness:.2f}, adoption: {adoption_rate:.2f}" + ) + + async def get_team_knowledge_stats(self, team_id: str) -> Dict[str, Any]: + """Get comprehensive team knowledge statistics""" + + async with self.pool.acquire() as conn: + # Basic statistics + basic_stats = await conn.fetchrow( + """ + SELECT + COUNT(*) as total_knowledge, + COUNT(DISTINCT knowledge_category) as categories, + COUNT(DISTINCT unnest(contributing_agents)) as contributing_agents, + AVG(team_relevance_score) as avg_relevance, + AVG(effectiveness_score) as avg_effectiveness, + AVG(agent_adoption_rate) as avg_adoption_rate + FROM team_knowledge_base + WHERE team_id = $1 + """, + team_id, + ) + + # Category breakdown + category_stats = await conn.fetch( + """ + SELECT + knowledge_category, + COUNT(*) as count, + AVG(effectiveness_score) as avg_effectiveness, + AVG(agent_adoption_rate) as avg_adoption + FROM team_knowledge_base + WHERE team_id = $1 + GROUP BY knowledge_category + ORDER BY count DESC + """, + team_id, + ) + + # Most effective knowledge + top_knowledge = await conn.fetch( + """ + SELECT + title, + knowledge_category, + effectiveness_score, + agent_adoption_rate + FROM team_knowledge_base + WHERE team_id = $1 + ORDER BY effectiveness_score DESC + LIMIT 5 + """, + team_id, + ) + + return { + "team_id": team_id, + "basic_stats": dict(basic_stats) if basic_stats else {}, + "category_breakdown": [dict(cat) for cat in category_stats], + "top_knowledge": [dict(know) for know in top_knowledge], + "generated_at": datetime.now().isoformat(), + } + + async def _search_team_specific_knowledge( + self, + conn, + team_id: str, + query_embedding: List[float], + categories: Optional[List[KnowledgeCategory]], + content_types: Optional[List[ContentType]], + limit: int, + min_similarity: float, + ) -> List[TeamKnowledgeSearchResult]: + """Search team-specific knowledge base""" + + # Build query conditions + where_conditions = ["team_id = $2"] + params = [query_embedding, team_id] + param_idx = 3 + + if categories: + where_conditions.append(f"knowledge_category = ANY(${param_idx})") + params.append([cat.value for cat in categories]) + param_idx += 1 + + if content_types: + where_conditions.append(f"content_type = ANY(${param_idx})") + params.append([ct.value for ct in content_types]) + param_idx += 1 + + where_conditions.append(f"(1 - (embedding <=> $1)) >= ${param_idx}") + params.append(min_similarity) + param_idx += 1 + + where_clause = "WHERE " + " AND ".join(where_conditions) + + results = await conn.fetch( + f""" + SELECT + *, + (1 - (embedding <=> $1)) as similarity_score + FROM team_knowledge_base + {where_clause} + ORDER BY similarity_score DESC, effectiveness_score DESC + LIMIT ${param_idx} + """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params + *params, + limit, + ) + + search_results = [] + for row in results: + team_knowledge = self._row_to_team_knowledge(row) + + # Calculate team fit score based on adoption and effectiveness + team_fit_score = ( + team_knowledge.agent_adoption_rate * 0.4 + + team_knowledge.effectiveness_score * 0.6 + ) + + combined_score = ( + float(row["similarity_score"]) * 0.4 + + team_knowledge.team_relevance_score * 0.3 + + team_fit_score * 0.3 + ) + + search_results.append( + TeamKnowledgeSearchResult( + team_knowledge=team_knowledge, + similarity_score=float(row["similarity_score"]), + relevance_score=team_knowledge.team_relevance_score, + team_fit_score=team_fit_score, + combined_score=combined_score, + ) + ) + + return search_results + + async def _search_org_knowledge_for_team( + self, + conn, + team_id: str, + query: str, + categories: Optional[List[KnowledgeCategory]], + content_types: Optional[List[ContentType]], + limit: int, + min_similarity: float, + ) -> List[TeamKnowledgeSearchResult]: + """Search organization knowledge filtered for team relevance""" + + # Get organization ID for the team + org_id = await conn.fetchval( + """ + SELECT organization_id FROM teams WHERE id = $1 + """, + team_id, + ) + + if not org_id: + return [] + + # Search organization knowledge + org_results = await self.org_rag_manager.search_knowledge( + organization_id=str(org_id), + query=query, + categories=categories, + content_types=content_types, + limit=limit * 2, # Get more to filter for team relevance + min_similarity=min_similarity, + requester_team_id=team_id, + ) + + # Convert to team knowledge search results with team relevance scoring + team_results = [] + for org_result in org_results: + # Calculate team relevance based on source and usage + team_relevance = await self._calculate_team_relevance( + conn, team_id, org_result.knowledge + ) + + if team_relevance >= self.min_team_relevance: + # Create pseudo team knowledge for consistent interface + pseudo_team_knowledge = TeamKnowledge( + id=org_result.knowledge.id, + team_id=team_id, + organization_id=org_result.knowledge.organization_id, + title=org_result.knowledge.title, + content=org_result.knowledge.content, + content_type=org_result.knowledge.content_type, + knowledge_category=org_result.knowledge.knowledge_category, + embedding=org_result.knowledge.embedding, + source_type=org_result.knowledge.source_type, + contributing_agents=[], + source_knowledge_ids=[org_result.knowledge.id], + aggregation_method="organization_filter", + team_relevance_score=team_relevance, + agent_adoption_rate=0.0, + effectiveness_score=org_result.knowledge.success_correlation, + visibility_level=org_result.knowledge.visibility_level, + metadata=org_result.knowledge.metadata, + tags=org_result.knowledge.tags, + created_at=org_result.knowledge.created_at, + updated_at=org_result.knowledge.updated_at, + last_accessed=org_result.knowledge.last_accessed, + ) + + combined_score = ( + org_result.similarity_score * 0.5 + team_relevance * 0.5 + ) + + team_results.append( + TeamKnowledgeSearchResult( + team_knowledge=pseudo_team_knowledge, + similarity_score=org_result.similarity_score, + relevance_score=org_result.relevance_score, + team_fit_score=team_relevance, + combined_score=combined_score, + ) + ) + + return team_results[:limit] + + async def _calculate_team_relevance( + self, conn, team_id: str, org_knowledge + ) -> float: + """Calculate how relevant organization knowledge is for a specific team""" + + relevance_factors = [] + + # Factor 1: Source team match + if org_knowledge.source_team_id == team_id: + relevance_factors.append(1.0) + elif org_knowledge.source_team_id: + # Check if source team is similar to current team + team_similarity = await self._calculate_team_similarity( + conn, team_id, org_knowledge.source_team_id + ) + relevance_factors.append(team_similarity) + else: + relevance_factors.append(0.3) # No team context + + # Factor 2: Category relevance to team's work + team_categories = await self._get_team_primary_categories(conn, team_id) + if org_knowledge.knowledge_category.value in team_categories: + relevance_factors.append(0.9) + else: + relevance_factors.append(0.4) + + # Factor 3: Usage by team agents + team_usage = ( + await conn.fetchval( + """ + SELECT COUNT(DISTINCT source_agent_id)::float / NULLIF( + (SELECT COUNT(*) FROM agents WHERE team_id = $1), 0 + ) + FROM organization_knowledge_base + WHERE id = $2 AND source_agent_id IN ( + SELECT id FROM agents WHERE team_id = $1 + ) + """, + team_id, + org_knowledge.id, + ) + or 0.0 + ) + relevance_factors.append(team_usage) + + # Factor 4: Base quality and relevance + relevance_factors.append(org_knowledge.quality_score) + relevance_factors.append(org_knowledge.relevance_score) + + # Calculate weighted average + weights = [0.3, 0.25, 0.25, 0.1, 0.1] + team_relevance = sum(f * w for f, w in zip(relevance_factors, weights)) + + return min(1.0, max(0.0, team_relevance)) + + async def _calculate_team_similarity( + self, conn, team_id1: str, team_id2: str + ) -> float: + """Calculate similarity between two teams based on their work patterns""" + + # Simple implementation based on team type and settings + team_info = await conn.fetch( + """ + SELECT id, team_type, settings FROM teams + WHERE id IN ($1, $2) + """, + team_id1, + team_id2, + ) + + if len(team_info) != 2: + return 0.0 + + team1, team2 = team_info + + # Type similarity + type_similarity = 1.0 if team1["team_type"] == team2["team_type"] else 0.5 + + # Settings similarity (simplified) + settings1 = team1["settings"] or {} + settings2 = team2["settings"] or {} + + common_keys = set(settings1.keys()) & set(settings2.keys()) + if common_keys: + settings_similarity = sum( + 1.0 if settings1.get(key) == settings2.get(key) else 0.0 + for key in common_keys + ) / len(common_keys) + else: + settings_similarity = 0.5 + + return type_similarity * 0.7 + settings_similarity * 0.3 + + async def _get_team_primary_categories(self, conn, team_id: str) -> List[str]: + """Get primary knowledge categories this team works with""" + + categories = await conn.fetch( + """ + SELECT knowledge_category, COUNT(*) as usage_count + FROM team_knowledge_base + WHERE team_id = $1 + GROUP BY knowledge_category + ORDER BY usage_count DESC + LIMIT 3 + """, + team_id, + ) + + return [cat["knowledge_category"] for cat in categories] + + def _generate_embedding(self, text: str) -> List[float]: + """Generate embedding for text using sentence transformers""" + try: + embedding = self.embedding_model.encode(text, convert_to_tensor=False) + return embedding.tolist() + except Exception as e: + logger.error(f"Error generating embedding: {e}") + return [0.0] * self.embedding_dim + + def _row_to_team_knowledge(self, row) -> TeamKnowledge: + """Convert database row to TeamKnowledge object""" + return TeamKnowledge( + id=str(row["id"]), + team_id=str(row["team_id"]), + organization_id=str(row["organization_id"]), + title=row["title"], + content=row["content"], + content_type=ContentType(row["content_type"]), + knowledge_category=KnowledgeCategory(row["knowledge_category"]), + embedding=row["embedding"] if row["embedding"] else None, + source_type=SourceType(row["source_type"]), + contributing_agents=row["contributing_agents"] or [], + source_knowledge_ids=row["source_knowledge_ids"] or [], + aggregation_method=row["aggregation_method"], + team_relevance_score=row["team_relevance_score"], + agent_adoption_rate=row["agent_adoption_rate"], + effectiveness_score=row["effectiveness_score"], + visibility_level=VisibilityLevel(row["visibility_level"]), + metadata=( + json.loads(row["metadata"]) + if isinstance(row["metadata"], str) + else row["metadata"] + ), + tags=row["tags"] or [], + created_at=row["created_at"], + updated_at=row["updated_at"], + last_accessed=row["last_accessed"], + ) + + def _build_context_query(self, task_context: Dict[str, Any]) -> str: + """Build a search query from task context""" + query_parts = [] + + if task_context.get("task_type"): + query_parts.append(task_context["task_type"]) + + if task_context.get("description"): + query_parts.append(task_context["description"]) + + if task_context.get("technologies"): + query_parts.extend(task_context["technologies"]) + + if task_context.get("domain"): + query_parts.append(task_context["domain"]) + + return " ".join(query_parts) + + async def _analyze_memories_for_aggregation( + self, agent_memories: List + ) -> Dict[str, Any]: + """Analyze agent memories to determine if they should be aggregated""" + + if not agent_memories: + return {"aggregation_value": 0.0} + + # Simple aggregation analysis + # In practice, this could use more sophisticated NLP + + # Calculate average confidence and success correlation + avg_confidence = sum(mem["confidence_score"] for mem in agent_memories) / len( + agent_memories + ) + avg_success = sum( + mem.get("success_correlation", 0.0) for mem in agent_memories + ) / len(agent_memories) + + # Find common themes + all_content = " ".join(mem["content"] for mem in agent_memories) + + # Determine primary category + categories = [mem.get("memory_type", "general") for mem in agent_memories] + primary_category = ( + max(set(categories), key=categories.count) if categories else "general" + ) + + # Create aggregated content (simplified) + title = f"Team Knowledge: {primary_category.replace('_', ' ').title()}" + content = ( + f"Aggregated knowledge from {len(agent_memories)} agent experiences:\n\n" + + all_content[:1000] + ) + + # Map memory type to knowledge category + category_mapping = { + "code_pattern": KnowledgeCategory.DEVELOPMENT, + "task_outcome": KnowledgeCategory.PROCESS, + "debugging": KnowledgeCategory.TROUBLESHOOTING, + "optimization": KnowledgeCategory.DEVELOPMENT, + "testing": KnowledgeCategory.TESTING, + } + + knowledge_category = category_mapping.get( + primary_category, KnowledgeCategory.DEVELOPMENT + ) + + # Determine content type + content_type = ( + ContentType.CODE if "code" in primary_category else ContentType.TEXT + ) + + # Calculate aggregation value + aggregation_value = min(1.0, (avg_confidence + avg_success) / 2.0) + + return { + "aggregation_value": aggregation_value, + "title": title, + "content": content, + "content_type": content_type, + "category": knowledge_category, + "metadata": { + "source_memory_count": len(agent_memories), + "avg_confidence": avg_confidence, + "avg_success_correlation": avg_success, + "primary_type": primary_category, + }, + "tags": [primary_category, "aggregated", "agent_contribution"], + } From 2099234791a77e36e872700ce0559942d93a1d63 Mon Sep 17 00:00:00 2001 From: Claude Code Date: Tue, 25 Aug 2026 08:42:58 +0000 Subject: [PATCH 3/3] Revert "chore(governance): reconcile managed files to FuzeSDLC v1 [skip ci]" This reverts commit b09b32a6897c37703c2ffb8fc95b84deb9db6c9a. --- .fuze/repo-manifest.schema.json | 40 +- hierarchy_endpoints.py | 862 +- .../orchestrator/agent_expertise_tracker.py | 1040 +- services/orchestrator/claude_code_wrapper.py | 1684 +-- services/orchestrator/claude_sdk_manager.py | 1004 +- .../context_enhancement_service.py | 1548 +- services/orchestrator/context_service.py | 304 +- services/orchestrator/conversation_manager.py | 1204 +- .../orchestrator/coordination_endpoints.py | 1250 +- .../orchestrator/goal_conversation_service.py | 2080 +-- services/orchestrator/hierarchy_endpoints.py | 786 +- .../knowledge_propagation_engine.py | 1916 +-- services/orchestrator/main.py | 12026 ++++++++-------- services/orchestrator/mcp_integration.py | 1318 +- services/orchestrator/model_configuration.py | 1116 +- .../orchestrator/multi_agent_coordinator.py | 1876 +-- .../orchestrator/task_execution_engine.py | 2612 ++-- .../orchestrator/task_knowledge_extractor.py | 1758 +-- services/orchestrator/task_queue.py | 248 +- .../orchestrator/team_knowledge_manager.py | 1738 +-- 20 files changed, 18224 insertions(+), 18186 deletions(-) diff --git a/.fuze/repo-manifest.schema.json b/.fuze/repo-manifest.schema.json index 3db2bbe..38496c7 100644 --- a/.fuze/repo-manifest.schema.json +++ b/.fuze/repo-manifest.schema.json @@ -231,8 +231,12 @@ "platformAuth": { "type": "object", "additionalProperties": false, - "description": "NEW BLOCK \u2014 no repo declares it yet, and that is the point: it gates the platform-auth capability. Consume @fuzefront/auth (published as @izzywdev/fuzefront-auth) rather than a bespoke verifier. A product NEVER calls Permit directly; it knows exactly one thing, the base URL of FuzeFront's Security API.", + "description": "NEW BLOCK. Consume @fuzefront/auth (published as @izzywdev/fuzefront-auth) rather than a bespoke verifier. A product NEVER calls Permit directly; it knows exactly one thing, the base URL of FuzeFront's Security API. gate-platform-auth ENFORCES BY DEFAULT \u2014 this block is how a repo opts OUT, not how it opts in.", "properties": { + "enforce": { + "type": "boolean", + "description": "Ratchet for gate-platform-auth, and it is OPT-OUT: absent means ENFORCING. Set false only to silence the gate while a repo migrates, and only together with `reason` \u2014 an `enforce: false` with no reason is ignored and the gate enforces anyway, because an undocumented opt-out is indistinguishable from an oversight. The earlier opt-in shape was chosen to avoid redding the fleet on pre-existing violations, but that is how gate-identifier reached zero adoption across 21 repos: a check nobody enabled is indistinguishable from a check that does not exist. What actually prevents a `|| true` is visibility, not coldness \u2014 an `enforce: false` naming a repo and a reason is greppable and countable; `|| true` in a workflow is neither." + }, "mode": { "enum": [ "federated-jwks", @@ -251,6 +255,10 @@ }, "note": { "type": "string" + }, + "reason": { + "type": "string", + "description": "REQUIRED when enforce is false. What blocks adoption and who owns closing it. This is the whole cost of the escape hatch: the opt-out must read as debt someone wrote down, not as a setting someone left alone." } } }, @@ -667,6 +675,36 @@ } } }, + "dataTier": { + "type": "array", + "description": "Declarative data-tier provisioning request (the IaC hand-off to FuzeInfra). FuzeInfra's reconciler consumes each entry: it ensures the per-service role exists AND is GRANTED the declared privileges on the declared database, then VERIFIES the role can actually read/write it (fail-loud if a role can auth but not access its DB). Replaces the old ad-hoc '@claude please provision' request (governance/shared-cluster-deploy.md §5). Every store the product's role authenticates to MUST be declared here, with the exact database name the app uses — a role granted on the wrong db name is the classic silent-empty-data bug.", + "items": { + "type": "object", + "additionalProperties": false, + "required": ["store", "database", "role"], + "properties": { + "store": { "enum": ["postgres", "mongo", "redis", "neo4j", "chroma"], "description": "Shared datastore this role needs access in." }, + "database": { "type": "string", "description": "The exact database/keyspace name the app reads/writes (e.g. robot_catalog). The role MUST be granted on THIS name; provisioning verifies it." }, + "role": { "type": "string", "description": "The per-service role/user (e.g. mendys)." }, + "privileges": { "enum": ["readWrite", "read", "admin"], "default": "readWrite", "description": "Privilege level to grant the role on `database`." }, + "authSource": { "type": "string", "description": "Mongo authSource db the role authenticates against (e.g. admin), when it differs from `database`." } + } + } + }, + "egress": { + "type": "array", + "description": "External hosts the product's pods need outbound HTTPS to. The shared cluster is egress-restricted (HTTP-only behind the Cloudflare tunnel; no default outbound to third-party APIs), so every external dependency MUST be declared here. FuzeInfra's reconciler turns these into namespace egress allow-rules (NetworkPolicy / egress gateway). Declare each third-party API explicitly (e.g. LLM providers).", + "items": { + "type": "object", + "additionalProperties": false, + "required": ["host"], + "properties": { + "host": { "type": "string", "description": "FQDN, e.g. api.openai.com." }, + "port": { "type": "integer", "default": 443, "description": "Destination port (default 443)." }, + "reason": { "type": "string", "description": "Why the product needs it (e.g. 'AI keyword generation')." } + } + } + }, "dependsOn": { "type": "array", "description": "Product-to-product dependencies this repo consumes beyond the spine (e.g. FuzeService dependsOn FuzeContact, FuzeBI).", diff --git a/hierarchy_endpoints.py b/hierarchy_endpoints.py index e48d494..acaa613 100644 --- a/hierarchy_endpoints.py +++ b/hierarchy_endpoints.py @@ -1,432 +1,432 @@ -#!/usr/bin/env python3 -""" -Simple FastAPI service that adds organization/team endpoints -and proxies other requests to the orchestrator -""" - -import asyncio -import asyncpg -import json -import uuid -import httpx -from contextlib import asynccontextmanager -from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect, Depends -from fastapi.middleware.cors import CORSMiddleware -from fastapi.responses import JSONResponse -from pydantic import BaseModel -from typing import List, Optional, Dict, Set - -# Configuration -import os -# SECURITY (issue #6): do not ship a real-looking DB password as a default. -DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://postgres:postgres@localhost:5434/ai_context") -ORCHESTRATOR_URL = os.getenv("ORCHESTRATOR_URL", "http://localhost:8000") - -try: - from auth import get_current_user, require_user, require_org_access, CurrentUser, authenticate_websocket -except Exception: # pragma: no cover - allow import from repo root or service dir - from services.orchestrator.auth import ( # type: ignore - get_current_user, require_user, require_org_access, CurrentUser, - authenticate_websocket, - ) - -@asynccontextmanager -async def lifespan(app: FastAPI): - """Application lifespan. - - Replaces the removed ``app.add_event_handler("startup"/"shutdown", ...)`` - API (dropped in Starlette 1.x). Behaviour is unchanged: open the asyncpg - pool on startup, close it on shutdown. ``startup``/``shutdown`` are - resolved at call time, so they may be defined further down the module. - """ - await startup() - try: - yield - finally: - await shutdown() - - -# SECURITY (issue #6 CRITICAL-1): authenticate every route by default (health -# and docs are on the allowlist inside get_current_user). -app = FastAPI( - title="FuzeAgent Hierarchy API", - version="1.0.0", - dependencies=[Depends(get_current_user)], - lifespan=lifespan, -) - -# SECURITY (issue #6 MEDIUM-2): explicit, non-wildcard origins when credentials -# are allowed (wildcard + credentials is both insecure and spec-invalid). -_cors_origins = [ - o.strip() - for o in os.getenv( - "CORS_ALLOW_ORIGINS", - "http://localhost:3000,http://localhost:3031,http://localhost", - ).split(",") - if o.strip() -] -app.add_middleware( - CORSMiddleware, - allow_origins=_cors_origins, - allow_credentials=True, - allow_methods=["*"], - allow_headers=["*"], -) - -# Database connection pool -db_pool = None - -# WebSocket connection manager -class ConnectionManager: - def __init__(self): - self.active_connections: Set[WebSocket] = set() - - async def connect(self, websocket: WebSocket): - await websocket.accept() - self.active_connections.add(websocket) - - def disconnect(self, websocket: WebSocket): - self.active_connections.discard(websocket) - - async def broadcast(self, message: dict): - disconnected = set() - for connection in self.active_connections: - try: - await connection.send_text(json.dumps(message)) - except: - disconnected.add(connection) - - # Remove disconnected clients - for connection in disconnected: - self.disconnect(connection) - -manager = ConnectionManager() - -async def startup(): - global db_pool - db_pool = await asyncpg.create_pool(DATABASE_URL, min_size=1, max_size=10) - -async def shutdown(): - if db_pool: - await db_pool.close() - -# NOTE: startup/shutdown are wired via the `lifespan` context manager defined -# above and passed to FastAPI(...); `add_event_handler` was removed in -# Starlette 1.x. - -# Pydantic models -class Organization(BaseModel): - id: str - name: str - description: Optional[str] = None - settings: dict = {} - created_at: str - updated_at: str - -class OrganizationCreate(BaseModel): - name: str - description: Optional[str] = None - settings: dict = {} - -class Team(BaseModel): - id: str - organization_id: str - name: str - description: Optional[str] = None - team_type: str = "general" - settings: dict = {} - created_at: str - updated_at: str - -class TeamCreate(BaseModel): - organization_id: str - name: str - description: Optional[str] = None - team_type: str = "general" - settings: dict = {} - -# Organization endpoints -@app.get("/organizations", response_model=List[Organization]) -async def get_organizations(): - async with db_pool.acquire() as conn: - rows = await conn.fetch(""" - SELECT - id::text, name, description, settings, - created_at::text, updated_at::text - FROM organizations - ORDER BY created_at DESC - """) - - organizations = [] - for row in rows: - organizations.append(Organization( - id=row['id'], - name=row['name'], - description=row['description'], - settings=json.loads(row['settings']) if row['settings'] else {}, - created_at=row['created_at'], - updated_at=row['updated_at'] - )) - - return organizations - -@app.post("/organizations", response_model=Organization) -async def create_organization(org_data: OrganizationCreate): - async with db_pool.acquire() as conn: - org_id = str(uuid.uuid4()) - row = await conn.fetchrow(""" - INSERT INTO organizations (id, name, description, settings) - VALUES ($1, $2, $3, $4) - RETURNING - id::text, name, description, settings, - created_at::text, updated_at::text - """, org_id, org_data.name, org_data.description, json.dumps(org_data.settings)) - - organization = Organization( - id=row['id'], - name=row['name'], - description=row['description'], - settings=json.loads(row['settings']) if row['settings'] else {}, - created_at=row['created_at'], - updated_at=row['updated_at'] - ) - - # Broadcast the change - await manager.broadcast({ - "type": "organization_created", - "data": organization.dict() - }) - - return organization - -@app.get("/organizations/{organization_id}", response_model=Organization) -async def get_organization( - organization_id: str, - user: CurrentUser = Depends(require_user), -): - # SECURITY (issue #6 HIGH-2 / BOLA): authorize the specific org id from the - # path; bare ``WHERE id = $1`` is not an authorization boundary. - require_org_access(organization_id, user) - async with db_pool.acquire() as conn: - row = await conn.fetchrow(""" - SELECT - id::text, name, description, settings, - created_at::text, updated_at::text - FROM organizations - WHERE id = $1 - """, organization_id) - - if not row: - raise HTTPException(status_code=404, detail="Organization not found") - - return Organization( - id=row['id'], - name=row['name'], - description=row['description'], - settings=json.loads(row['settings']) if row['settings'] else {}, - created_at=row['created_at'], - updated_at=row['updated_at'] - ) - -# Team endpoints -@app.get("/teams", response_model=List[Team]) -async def get_teams(organization_id: Optional[str] = None): - async with db_pool.acquire() as conn: - if organization_id: - rows = await conn.fetch(""" - SELECT - id::text, organization_id::text, name, description, - team_type, settings, created_at::text, updated_at::text - FROM teams - WHERE organization_id = $1 - ORDER BY created_at DESC - """, organization_id) - else: - rows = await conn.fetch(""" - SELECT - id::text, organization_id::text, name, description, - team_type, settings, created_at::text, updated_at::text - FROM teams - ORDER BY created_at DESC - """) - - teams = [] - for row in rows: - teams.append(Team( - id=row['id'], - organization_id=row['organization_id'], - name=row['name'], - description=row['description'], - team_type=row['team_type'], - settings=json.loads(row['settings']) if row['settings'] else {}, - created_at=row['created_at'], - updated_at=row['updated_at'] - )) - - return teams - -@app.post("/teams", response_model=Team) -async def create_team(team_data: TeamCreate): - async with db_pool.acquire() as conn: - # Verify organization exists - org_exists = await conn.fetchval( - "SELECT EXISTS(SELECT 1 FROM organizations WHERE id = $1)", - team_data.organization_id - ) - if not org_exists: - raise HTTPException(status_code=404, detail="Organization not found") - - team_id = str(uuid.uuid4()) - row = await conn.fetchrow(""" - INSERT INTO teams (id, organization_id, name, description, team_type, settings) - VALUES ($1, $2, $3, $4, $5, $6) - RETURNING - id::text, organization_id::text, name, description, - team_type, settings, created_at::text, updated_at::text - """, team_id, team_data.organization_id, team_data.name, - team_data.description, team_data.team_type, json.dumps(team_data.settings)) - - team = Team( - id=row['id'], - organization_id=row['organization_id'], - name=row['name'], - description=row['description'], - team_type=row['team_type'], - settings=json.loads(row['settings']) if row['settings'] else {}, - created_at=row['created_at'], - updated_at=row['updated_at'] - ) - - # Broadcast the change - await manager.broadcast({ - "type": "team_created", - "data": team.dict() - }) - - return team - -@app.get("/teams/{team_id}", response_model=Team) -async def get_team( - team_id: str, - user: CurrentUser = Depends(require_user), -): - async with db_pool.acquire() as conn: - row = await conn.fetchrow(""" - SELECT - id::text, organization_id::text, name, description, - team_type, settings, created_at::text, updated_at::text - FROM teams - WHERE id = $1 - """, team_id) - - if not row: - raise HTTPException(status_code=404, detail="Team not found") - # SECURITY (issue #6 HIGH-2): authorize via the team's parent org. - require_org_access(row['organization_id'], user) - return Team( - id=row['id'], - organization_id=row['organization_id'], - name=row['name'], - description=row['description'], - team_type=row['team_type'], - settings=json.loads(row['settings']) if row['settings'] else {}, - created_at=row['created_at'], - updated_at=row['updated_at'] - ) - -# WebSocket endpoint for real-time updates -@app.websocket("/ws") -async def websocket_endpoint(websocket: WebSocket): - # SECURITY (issue #6 CRITICAL-2): authenticate BEFORE accept(). The - # app-wide ``dependencies=[Depends(get_current_user)]`` is a no-op on - # WebSocket routes (WS handshakes have no HTTP response channel); every WS - # handler must call authenticate_websocket() first — matching the pattern - # already used by all handlers in services/orchestrator/main.py. - user = await authenticate_websocket(websocket) - if user is None: - return # authenticate_websocket already closed the socket (1008) - await manager.connect(websocket) - try: - while True: - # Keep the connection alive and listen for client pings - data = await websocket.receive_text() - if data == "ping": - await websocket.send_text("pong") - except WebSocketDisconnect: - manager.disconnect(websocket) - -# Proxy all other requests to the original orchestrator -@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"]) -async def proxy_to_orchestrator(path: str, request: Request): - url = f"{ORCHESTRATOR_URL}/{path}" - - # Handle CORS preflight requests - if request.method == "OPTIONS": - return JSONResponse( - content={}, - headers={ - "Access-Control-Allow-Origin": "*", - "Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, PATCH, OPTIONS", - "Access-Control-Allow-Headers": "*", - } - ) - - try: - async with httpx.AsyncClient(timeout=30.0) as client: - # Get request body if present - body = None - if request.method in ["POST", "PUT", "PATCH"]: - body = await request.body() - # Log the request for debugging - if path == "agents/from-template": - print(f"DEBUG: Proxying agents/from-template request") - print(f"DEBUG: URL: {url}") - print(f"DEBUG: Body: {body.decode() if body else 'None'}") - print(f"DEBUG: Headers: {dict(request.headers)}") - - # Forward the request - response = await client.request( - method=request.method, - url=url, - params=request.query_params, - content=body, - headers={k: v for k, v in request.headers.items() - if k.lower() not in ['host', 'content-length']}, - ) - - # Check if this was an agent creation and broadcast the change - if (request.method == "POST" and - (path == "agents" or path == "agents/from-template") and - response.status_code in [200, 201] and - response.headers.get("content-type", "").startswith("application/json")): - try: - agent_data = response.json() - await manager.broadcast({ - "type": "agent_created", - "data": agent_data - }) - except: - pass # Ignore broadcast errors - - return JSONResponse( - content=response.json() if response.headers.get("content-type", "").startswith("application/json") else response.text, - status_code=response.status_code, - headers={ - "Access-Control-Allow-Origin": "*", - **{k: v for k, v in response.headers.items() - if k.lower() not in ['content-length', 'transfer-encoding', 'connection']} - } - ) - - except httpx.RequestError as e: - print(f"ERROR: RequestError in proxy: {str(e)}") - raise HTTPException(status_code=503, detail=f"Failed to connect to orchestrator: {str(e)}") - except Exception as e: - print(f"ERROR: Exception in proxy: {str(e)}") - import traceback - traceback.print_exc() - raise HTTPException(status_code=500, detail=f"Proxy error: {str(e)}") - -if __name__ == "__main__": - import uvicorn +#!/usr/bin/env python3 +""" +Simple FastAPI service that adds organization/team endpoints +and proxies other requests to the orchestrator +""" + +import asyncio +import asyncpg +import json +import uuid +import httpx +from contextlib import asynccontextmanager +from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect, Depends +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import JSONResponse +from pydantic import BaseModel +from typing import List, Optional, Dict, Set + +# Configuration +import os +# SECURITY (issue #6): do not ship a real-looking DB password as a default. +DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://postgres:postgres@localhost:5434/ai_context") +ORCHESTRATOR_URL = os.getenv("ORCHESTRATOR_URL", "http://localhost:8000") + +try: + from auth import get_current_user, require_user, require_org_access, CurrentUser, authenticate_websocket +except Exception: # pragma: no cover - allow import from repo root or service dir + from services.orchestrator.auth import ( # type: ignore + get_current_user, require_user, require_org_access, CurrentUser, + authenticate_websocket, + ) + +@asynccontextmanager +async def lifespan(app: FastAPI): + """Application lifespan. + + Replaces the removed ``app.add_event_handler("startup"/"shutdown", ...)`` + API (dropped in Starlette 1.x). Behaviour is unchanged: open the asyncpg + pool on startup, close it on shutdown. ``startup``/``shutdown`` are + resolved at call time, so they may be defined further down the module. + """ + await startup() + try: + yield + finally: + await shutdown() + + +# SECURITY (issue #6 CRITICAL-1): authenticate every route by default (health +# and docs are on the allowlist inside get_current_user). +app = FastAPI( + title="FuzeAgent Hierarchy API", + version="1.0.0", + dependencies=[Depends(get_current_user)], + lifespan=lifespan, +) + +# SECURITY (issue #6 MEDIUM-2): explicit, non-wildcard origins when credentials +# are allowed (wildcard + credentials is both insecure and spec-invalid). +_cors_origins = [ + o.strip() + for o in os.getenv( + "CORS_ALLOW_ORIGINS", + "http://localhost:3000,http://localhost:3031,http://localhost", + ).split(",") + if o.strip() +] +app.add_middleware( + CORSMiddleware, + allow_origins=_cors_origins, + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +# Database connection pool +db_pool = None + +# WebSocket connection manager +class ConnectionManager: + def __init__(self): + self.active_connections: Set[WebSocket] = set() + + async def connect(self, websocket: WebSocket): + await websocket.accept() + self.active_connections.add(websocket) + + def disconnect(self, websocket: WebSocket): + self.active_connections.discard(websocket) + + async def broadcast(self, message: dict): + disconnected = set() + for connection in self.active_connections: + try: + await connection.send_text(json.dumps(message)) + except: + disconnected.add(connection) + + # Remove disconnected clients + for connection in disconnected: + self.disconnect(connection) + +manager = ConnectionManager() + +async def startup(): + global db_pool + db_pool = await asyncpg.create_pool(DATABASE_URL, min_size=1, max_size=10) + +async def shutdown(): + if db_pool: + await db_pool.close() + +# NOTE: startup/shutdown are wired via the `lifespan` context manager defined +# above and passed to FastAPI(...); `add_event_handler` was removed in +# Starlette 1.x. + +# Pydantic models +class Organization(BaseModel): + id: str + name: str + description: Optional[str] = None + settings: dict = {} + created_at: str + updated_at: str + +class OrganizationCreate(BaseModel): + name: str + description: Optional[str] = None + settings: dict = {} + +class Team(BaseModel): + id: str + organization_id: str + name: str + description: Optional[str] = None + team_type: str = "general" + settings: dict = {} + created_at: str + updated_at: str + +class TeamCreate(BaseModel): + organization_id: str + name: str + description: Optional[str] = None + team_type: str = "general" + settings: dict = {} + +# Organization endpoints +@app.get("/organizations", response_model=List[Organization]) +async def get_organizations(): + async with db_pool.acquire() as conn: + rows = await conn.fetch(""" + SELECT + id::text, name, description, settings, + created_at::text, updated_at::text + FROM organizations + ORDER BY created_at DESC + """) + + organizations = [] + for row in rows: + organizations.append(Organization( + id=row['id'], + name=row['name'], + description=row['description'], + settings=json.loads(row['settings']) if row['settings'] else {}, + created_at=row['created_at'], + updated_at=row['updated_at'] + )) + + return organizations + +@app.post("/organizations", response_model=Organization) +async def create_organization(org_data: OrganizationCreate): + async with db_pool.acquire() as conn: + org_id = str(uuid.uuid4()) + row = await conn.fetchrow(""" + INSERT INTO organizations (id, name, description, settings) + VALUES ($1, $2, $3, $4) + RETURNING + id::text, name, description, settings, + created_at::text, updated_at::text + """, org_id, org_data.name, org_data.description, json.dumps(org_data.settings)) + + organization = Organization( + id=row['id'], + name=row['name'], + description=row['description'], + settings=json.loads(row['settings']) if row['settings'] else {}, + created_at=row['created_at'], + updated_at=row['updated_at'] + ) + + # Broadcast the change + await manager.broadcast({ + "type": "organization_created", + "data": organization.dict() + }) + + return organization + +@app.get("/organizations/{organization_id}", response_model=Organization) +async def get_organization( + organization_id: str, + user: CurrentUser = Depends(require_user), +): + # SECURITY (issue #6 HIGH-2 / BOLA): authorize the specific org id from the + # path; bare ``WHERE id = $1`` is not an authorization boundary. + require_org_access(organization_id, user) + async with db_pool.acquire() as conn: + row = await conn.fetchrow(""" + SELECT + id::text, name, description, settings, + created_at::text, updated_at::text + FROM organizations + WHERE id = $1 + """, organization_id) + + if not row: + raise HTTPException(status_code=404, detail="Organization not found") + + return Organization( + id=row['id'], + name=row['name'], + description=row['description'], + settings=json.loads(row['settings']) if row['settings'] else {}, + created_at=row['created_at'], + updated_at=row['updated_at'] + ) + +# Team endpoints +@app.get("/teams", response_model=List[Team]) +async def get_teams(organization_id: Optional[str] = None): + async with db_pool.acquire() as conn: + if organization_id: + rows = await conn.fetch(""" + SELECT + id::text, organization_id::text, name, description, + team_type, settings, created_at::text, updated_at::text + FROM teams + WHERE organization_id = $1 + ORDER BY created_at DESC + """, organization_id) + else: + rows = await conn.fetch(""" + SELECT + id::text, organization_id::text, name, description, + team_type, settings, created_at::text, updated_at::text + FROM teams + ORDER BY created_at DESC + """) + + teams = [] + for row in rows: + teams.append(Team( + id=row['id'], + organization_id=row['organization_id'], + name=row['name'], + description=row['description'], + team_type=row['team_type'], + settings=json.loads(row['settings']) if row['settings'] else {}, + created_at=row['created_at'], + updated_at=row['updated_at'] + )) + + return teams + +@app.post("/teams", response_model=Team) +async def create_team(team_data: TeamCreate): + async with db_pool.acquire() as conn: + # Verify organization exists + org_exists = await conn.fetchval( + "SELECT EXISTS(SELECT 1 FROM organizations WHERE id = $1)", + team_data.organization_id + ) + if not org_exists: + raise HTTPException(status_code=404, detail="Organization not found") + + team_id = str(uuid.uuid4()) + row = await conn.fetchrow(""" + INSERT INTO teams (id, organization_id, name, description, team_type, settings) + VALUES ($1, $2, $3, $4, $5, $6) + RETURNING + id::text, organization_id::text, name, description, + team_type, settings, created_at::text, updated_at::text + """, team_id, team_data.organization_id, team_data.name, + team_data.description, team_data.team_type, json.dumps(team_data.settings)) + + team = Team( + id=row['id'], + organization_id=row['organization_id'], + name=row['name'], + description=row['description'], + team_type=row['team_type'], + settings=json.loads(row['settings']) if row['settings'] else {}, + created_at=row['created_at'], + updated_at=row['updated_at'] + ) + + # Broadcast the change + await manager.broadcast({ + "type": "team_created", + "data": team.dict() + }) + + return team + +@app.get("/teams/{team_id}", response_model=Team) +async def get_team( + team_id: str, + user: CurrentUser = Depends(require_user), +): + async with db_pool.acquire() as conn: + row = await conn.fetchrow(""" + SELECT + id::text, organization_id::text, name, description, + team_type, settings, created_at::text, updated_at::text + FROM teams + WHERE id = $1 + """, team_id) + + if not row: + raise HTTPException(status_code=404, detail="Team not found") + # SECURITY (issue #6 HIGH-2): authorize via the team's parent org. + require_org_access(row['organization_id'], user) + return Team( + id=row['id'], + organization_id=row['organization_id'], + name=row['name'], + description=row['description'], + team_type=row['team_type'], + settings=json.loads(row['settings']) if row['settings'] else {}, + created_at=row['created_at'], + updated_at=row['updated_at'] + ) + +# WebSocket endpoint for real-time updates +@app.websocket("/ws") +async def websocket_endpoint(websocket: WebSocket): + # SECURITY (issue #6 CRITICAL-2): authenticate BEFORE accept(). The + # app-wide ``dependencies=[Depends(get_current_user)]`` is a no-op on + # WebSocket routes (WS handshakes have no HTTP response channel); every WS + # handler must call authenticate_websocket() first — matching the pattern + # already used by all handlers in services/orchestrator/main.py. + user = await authenticate_websocket(websocket) + if user is None: + return # authenticate_websocket already closed the socket (1008) + await manager.connect(websocket) + try: + while True: + # Keep the connection alive and listen for client pings + data = await websocket.receive_text() + if data == "ping": + await websocket.send_text("pong") + except WebSocketDisconnect: + manager.disconnect(websocket) + +# Proxy all other requests to the original orchestrator +@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"]) +async def proxy_to_orchestrator(path: str, request: Request): + url = f"{ORCHESTRATOR_URL}/{path}" + + # Handle CORS preflight requests + if request.method == "OPTIONS": + return JSONResponse( + content={}, + headers={ + "Access-Control-Allow-Origin": "*", + "Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, PATCH, OPTIONS", + "Access-Control-Allow-Headers": "*", + } + ) + + try: + async with httpx.AsyncClient(timeout=30.0) as client: + # Get request body if present + body = None + if request.method in ["POST", "PUT", "PATCH"]: + body = await request.body() + # Log the request for debugging + if path == "agents/from-template": + print(f"DEBUG: Proxying agents/from-template request") + print(f"DEBUG: URL: {url}") + print(f"DEBUG: Body: {body.decode() if body else 'None'}") + print(f"DEBUG: Headers: {dict(request.headers)}") + + # Forward the request + response = await client.request( + method=request.method, + url=url, + params=request.query_params, + content=body, + headers={k: v for k, v in request.headers.items() + if k.lower() not in ['host', 'content-length']}, + ) + + # Check if this was an agent creation and broadcast the change + if (request.method == "POST" and + (path == "agents" or path == "agents/from-template") and + response.status_code in [200, 201] and + response.headers.get("content-type", "").startswith("application/json")): + try: + agent_data = response.json() + await manager.broadcast({ + "type": "agent_created", + "data": agent_data + }) + except: + pass # Ignore broadcast errors + + return JSONResponse( + content=response.json() if response.headers.get("content-type", "").startswith("application/json") else response.text, + status_code=response.status_code, + headers={ + "Access-Control-Allow-Origin": "*", + **{k: v for k, v in response.headers.items() + if k.lower() not in ['content-length', 'transfer-encoding', 'connection']} + } + ) + + except httpx.RequestError as e: + print(f"ERROR: RequestError in proxy: {str(e)}") + raise HTTPException(status_code=503, detail=f"Failed to connect to orchestrator: {str(e)}") + except Exception as e: + print(f"ERROR: Exception in proxy: {str(e)}") + import traceback + traceback.print_exc() + raise HTTPException(status_code=500, detail=f"Proxy error: {str(e)}") + +if __name__ == "__main__": + import uvicorn uvicorn.run(app, host="0.0.0.0", port=8006) \ No newline at end of file diff --git a/services/orchestrator/agent_expertise_tracker.py b/services/orchestrator/agent_expertise_tracker.py index 33009c2..1848811 100644 --- a/services/orchestrator/agent_expertise_tracker.py +++ b/services/orchestrator/agent_expertise_tracker.py @@ -1,520 +1,520 @@ -""" -Agent Expertise Tracker - -Provides analytics and insights into agent performance, learning patterns, -and expertise development across the FuzeAgent system. -""" - -import asyncio -import json -import logging -from dataclasses import dataclass -from datetime import datetime, timedelta -from typing import Any, Dict, List, Optional - -from .database import get_db_connection - -logger = logging.getLogger(__name__) - - -@dataclass -class ExpertiseInsight: - """Insight about agent expertise development""" - - agent_id: str - skill_area: str - insight_type: str # 'improving', 'declining', 'plateau', 'breakthrough' - description: str - confidence: float - evidence: Dict[str, Any] - timestamp: datetime - - -@dataclass -class AgentPerformanceMetrics: - """Performance metrics for an agent""" - - agent_id: str - total_tasks: int - success_rate: float - avg_expertise_level: float - improving_skills_count: int - declining_skills_count: int - memory_usage_stats: Dict[str, Any] - recent_performance_trend: str - top_skill_areas: List[Dict[str, Any]] - - -class AgentExpertiseTracker: - """ - Tracks and analyzes agent expertise development, providing insights - into learning patterns, performance trends, and optimization opportunities. - """ - - def __init__(self, database_url: str): - self.database_url = database_url - self.insights_cache: Dict[str, List[ExpertiseInsight]] = {} - self.metrics_cache: Dict[str, AgentPerformanceMetrics] = {} - self.cache_ttl = 300 # 5 minutes - self.last_cache_update = {} - - async def get_agent_performance_metrics( - self, agent_id: str - ) -> Optional[AgentPerformanceMetrics]: - """Get comprehensive performance metrics for an agent""" - - # Check cache first - if ( - agent_id in self.metrics_cache - and agent_id in self.last_cache_update - and (datetime.now() - self.last_cache_update[agent_id]).total_seconds() - < self.cache_ttl - ): - return self.metrics_cache[agent_id] - - try: - async with get_db_connection() as conn: - # Get basic performance stats - basic_stats = await conn.fetchrow( - """ - SELECT - COUNT(DISTINCT am.task_id) as total_tasks, - AVG(CASE WHEN am.memory_type = 'success' THEN 1.0 ELSE 0.0 END) as success_rate, - COUNT(DISTINCT am.id) as total_memories - FROM agent_memory am - WHERE am.agent_id = $1 - """, - agent_id, - ) - - # Get expertise summary - expertise_stats = await conn.fetchrow( - """ - SELECT - AVG(expertise_level) as avg_expertise_level, - COUNT(CASE WHEN performance_trend = 'improving' THEN 1 END) as improving_skills, - COUNT(CASE WHEN performance_trend = 'declining' THEN 1 END) as declining_skills - FROM agent_expertise - WHERE agent_id = $1 - """, - agent_id, - ) - - # Get memory usage statistics - memory_stats = await conn.fetchrow( - """ - SELECT - COUNT(*) as total_memories, - AVG(confidence_score) as avg_confidence, - SUM(usage_count) as total_usage, - COUNT(DISTINCT memory_type) as memory_types_used - FROM agent_memory - WHERE agent_id = $1 - """, - agent_id, - ) - - # Get top skill areas - top_skills = await conn.fetch( - """ - SELECT skill_area, expertise_level, success_rate, task_count, performance_trend - FROM agent_expertise - WHERE agent_id = $1 - ORDER BY expertise_level DESC, success_rate DESC - LIMIT 5 - """, - agent_id, - ) - - # Determine recent performance trend - recent_trend = await self._calculate_recent_trend(agent_id, conn) - - # Build metrics object - metrics = AgentPerformanceMetrics( - agent_id=agent_id, - total_tasks=basic_stats["total_tasks"] or 0, - success_rate=basic_stats["success_rate"] or 0.0, - avg_expertise_level=expertise_stats["avg_expertise_level"] or 0.0, - improving_skills_count=expertise_stats["improving_skills"] or 0, - declining_skills_count=expertise_stats["declining_skills"] or 0, - memory_usage_stats={ - "total_memories": memory_stats["total_memories"] or 0, - "avg_confidence": ( - float(memory_stats["avg_confidence"]) - if memory_stats["avg_confidence"] - else 0.0 - ), - "total_usage": memory_stats["total_usage"] or 0, - "memory_types_used": memory_stats["memory_types_used"] or 0, - }, - recent_performance_trend=recent_trend, - top_skill_areas=[dict(skill) for skill in top_skills], - ) - - # Cache the result - self.metrics_cache[agent_id] = metrics - self.last_cache_update[agent_id] = datetime.now() - - return metrics - - except Exception as e: - logger.error(f"Error getting performance metrics for agent {agent_id}: {e}") - return None - - async def generate_expertise_insights( - self, agent_id: str - ) -> List[ExpertiseInsight]: - """Generate insights about agent expertise development""" - - # Check cache first - if ( - agent_id in self.insights_cache - and agent_id in self.last_cache_update - and (datetime.now() - self.last_cache_update[agent_id]).total_seconds() - < self.cache_ttl - ): - return self.insights_cache[agent_id] - - insights = [] - - try: - async with get_db_connection() as conn: - # Analyze learning velocity patterns - learning_insights = await self._analyze_learning_velocity( - agent_id, conn - ) - insights.extend(learning_insights) - - # Analyze skill development patterns - skill_insights = await self._analyze_skill_development(agent_id, conn) - insights.extend(skill_insights) - - # Analyze memory usage patterns - memory_insights = await self._analyze_memory_patterns(agent_id, conn) - insights.extend(memory_insights) - - # Cache the results - self.insights_cache[agent_id] = insights - self.last_cache_update[agent_id] = datetime.now() - - except Exception as e: - logger.error(f"Error generating insights for agent {agent_id}: {e}") - - return insights - - async def get_system_wide_expertise_summary(self) -> Dict[str, Any]: - """Get system-wide expertise and performance summary""" - - try: - async with get_db_connection() as conn: - # Overall system stats - system_stats = await conn.fetchrow(""" - SELECT - COUNT(DISTINCT a.id) as total_agents, - COUNT(DISTINCT ae.skill_area) as total_skill_areas, - AVG(ae.expertise_level) as avg_system_expertise, - COUNT(CASE WHEN ae.performance_trend = 'improving' THEN 1 END) as improving_agents, - COUNT(CASE WHEN ae.performance_trend = 'declining' THEN 1 END) as declining_agents - FROM agents a - LEFT JOIN agent_expertise ae ON a.id = ae.agent_id - """) - - # Memory system stats - memory_stats = await conn.fetchrow(""" - SELECT - COUNT(*) as total_memories, - AVG(confidence_score) as avg_confidence, - SUM(usage_count) as total_usage, - COUNT(DISTINCT agent_id) as agents_with_memory - FROM agent_memory - """) - - # Top performing skill areas - top_skill_areas = await conn.fetch(""" - SELECT - skill_area, - COUNT(*) as agent_count, - AVG(expertise_level) as avg_expertise, - AVG(success_rate) as avg_success_rate - FROM agent_expertise - GROUP BY skill_area - ORDER BY avg_expertise DESC, avg_success_rate DESC - LIMIT 10 - """) - - # Recent activity - recent_activity = await conn.fetchrow(""" - SELECT - COUNT(CASE WHEN created_at > NOW() - INTERVAL '24 hours' THEN 1 END) as memories_24h, - COUNT(CASE WHEN created_at > NOW() - INTERVAL '7 days' THEN 1 END) as memories_7d, - COUNT(DISTINCT CASE WHEN created_at > NOW() - INTERVAL '24 hours' THEN agent_id END) as active_agents_24h - FROM agent_memory - """) - - return { - "system_stats": dict(system_stats) if system_stats else {}, - "memory_stats": dict(memory_stats) if memory_stats else {}, - "top_skill_areas": [dict(skill) for skill in top_skill_areas], - "recent_activity": dict(recent_activity) if recent_activity else {}, - "timestamp": datetime.now().isoformat(), - } - - except Exception as e: - logger.error(f"Error getting system-wide expertise summary: {e}") - return {"error": str(e)} - - async def _calculate_recent_trend(self, agent_id: str, conn) -> str: - """Calculate recent performance trend for an agent""" - - try: - # Get recent task outcomes - recent_outcomes = await conn.fetch( - """ - SELECT - DATE_TRUNC('day', created_at) as date, - AVG(CASE WHEN memory_type = 'success' THEN 1.0 ELSE 0.0 END) as daily_success_rate - FROM agent_memory - WHERE agent_id = $1 - AND created_at > NOW() - INTERVAL '14 days' - AND memory_type IN ('success', 'task_outcome') - GROUP BY DATE_TRUNC('day', created_at) - ORDER BY date DESC - LIMIT 7 - """, - agent_id, - ) - - if len(recent_outcomes) < 3: - return "insufficient_data" - - # Calculate trend - success_rates = [ - float(row["daily_success_rate"]) for row in recent_outcomes - ] - - # Simple linear trend calculation - if len(success_rates) >= 3: - early_avg = sum(success_rates[-3:]) / 3 - recent_avg = sum(success_rates[:3]) / 3 - - if recent_avg > early_avg + 0.1: - return "improving" - elif recent_avg < early_avg - 0.1: - return "declining" - else: - return "stable" - - return "stable" - - except Exception as e: - logger.error(f"Error calculating recent trend: {e}") - return "unknown" - - async def _analyze_learning_velocity( - self, agent_id: str, conn - ) -> List[ExpertiseInsight]: - """Analyze learning velocity patterns""" - - insights = [] - - try: - # Get skills with high learning velocity - fast_learners = await conn.fetch( - """ - SELECT skill_area, learning_velocity, expertise_level, task_count - FROM agent_expertise - WHERE agent_id = $1 AND learning_velocity > 0.1 - ORDER BY learning_velocity DESC - """, - agent_id, - ) - - for skill in fast_learners: - insights.append( - ExpertiseInsight( - agent_id=agent_id, - skill_area=skill["skill_area"], - insight_type="improving", - description=f"Rapid improvement in {skill['skill_area']} with velocity {skill['learning_velocity']:.2f}", - confidence=0.8, - evidence={ - "learning_velocity": float(skill["learning_velocity"]), - "expertise_level": float(skill["expertise_level"]), - "task_count": skill["task_count"], - }, - timestamp=datetime.now(), - ) - ) - - # Get skills with declining performance - declining_skills = await conn.fetch( - """ - SELECT skill_area, learning_velocity, expertise_level, task_count - FROM agent_expertise - WHERE agent_id = $1 AND learning_velocity < -0.05 - ORDER BY learning_velocity ASC - """, - agent_id, - ) - - for skill in declining_skills: - insights.append( - ExpertiseInsight( - agent_id=agent_id, - skill_area=skill["skill_area"], - insight_type="declining", - description=f"Performance decline in {skill['skill_area']} - may need attention", - confidence=0.7, - evidence={ - "learning_velocity": float(skill["learning_velocity"]), - "expertise_level": float(skill["expertise_level"]), - "task_count": skill["task_count"], - }, - timestamp=datetime.now(), - ) - ) - - except Exception as e: - logger.error(f"Error analyzing learning velocity: {e}") - - return insights - - async def _analyze_skill_development( - self, agent_id: str, conn - ) -> List[ExpertiseInsight]: - """Analyze skill development patterns""" - - insights = [] - - try: - # Find breakthrough moments (significant expertise jumps) - breakthroughs = await conn.fetch( - """ - SELECT skill_area, expertise_level, success_rate, task_count - FROM agent_expertise - WHERE agent_id = $1 - AND expertise_level > 0.7 - AND success_rate > 0.8 - AND task_count >= 5 - """, - agent_id, - ) - - for breakthrough in breakthroughs: - insights.append( - ExpertiseInsight( - agent_id=agent_id, - skill_area=breakthrough["skill_area"], - insight_type="breakthrough", - description=f"Expert level achieved in {breakthrough['skill_area']} with {breakthrough['success_rate']:.1%} success rate", - confidence=0.9, - evidence={ - "expertise_level": float(breakthrough["expertise_level"]), - "success_rate": float(breakthrough["success_rate"]), - "task_count": breakthrough["task_count"], - }, - timestamp=datetime.now(), - ) - ) - - # Find plateau situations (high task count but low expertise) - plateaus = await conn.fetch( - """ - SELECT skill_area, expertise_level, success_rate, task_count - FROM agent_expertise - WHERE agent_id = $1 - AND task_count > 10 - AND expertise_level < 0.4 - AND learning_velocity BETWEEN -0.02 AND 0.02 - """, - agent_id, - ) - - for plateau in plateaus: - insights.append( - ExpertiseInsight( - agent_id=agent_id, - skill_area=plateau["skill_area"], - insight_type="plateau", - description=f"Learning plateau in {plateau['skill_area']} - consider new approaches", - confidence=0.6, - evidence={ - "expertise_level": float(plateau["expertise_level"]), - "success_rate": float(plateau["success_rate"]), - "task_count": plateau["task_count"], - }, - timestamp=datetime.now(), - ) - ) - - except Exception as e: - logger.error(f"Error analyzing skill development: {e}") - - return insights - - async def _analyze_memory_patterns( - self, agent_id: str, conn - ) -> List[ExpertiseInsight]: - """Analyze memory usage and effectiveness patterns""" - - insights = [] - - try: - # Analyze memory types and their effectiveness - memory_effectiveness = await conn.fetchrow( - """ - SELECT - COUNT(*) as total_memories, - AVG(usage_count) as avg_usage, - AVG(confidence_score) as avg_confidence, - COUNT(CASE WHEN usage_count > 5 THEN 1 END) as high_usage_memories - FROM agent_memory - WHERE agent_id = $1 - """, - agent_id, - ) - - if memory_effectiveness and memory_effectiveness["total_memories"] > 50: - high_usage_ratio = ( - memory_effectiveness["high_usage_memories"] - / memory_effectiveness["total_memories"] - ) - - if ( - high_usage_ratio > 0.2 - ): # More than 20% of memories are highly reused - insights.append( - ExpertiseInsight( - agent_id=agent_id, - skill_area="memory_management", - insight_type="improving", - description=f"Excellent memory reuse patterns - {high_usage_ratio:.1%} of memories are frequently accessed", - confidence=0.8, - evidence={ - "total_memories": memory_effectiveness[ - "total_memories" - ], - "high_usage_ratio": high_usage_ratio, - "avg_confidence": float( - memory_effectiveness["avg_confidence"] - ), - }, - timestamp=datetime.now(), - ) - ) - - except Exception as e: - logger.error(f"Error analyzing memory patterns: {e}") - - return insights - - async def clear_cache(self, agent_id: Optional[str] = None): - """Clear analytics cache""" - if agent_id: - self.insights_cache.pop(agent_id, None) - self.metrics_cache.pop(agent_id, None) - self.last_cache_update.pop(agent_id, None) - else: - self.insights_cache.clear() - self.metrics_cache.clear() - self.last_cache_update.clear() +""" +Agent Expertise Tracker + +Provides analytics and insights into agent performance, learning patterns, +and expertise development across the FuzeAgent system. +""" + +import asyncio +import json +import logging +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Any, Dict, List, Optional + +from .database import get_db_connection + +logger = logging.getLogger(__name__) + + +@dataclass +class ExpertiseInsight: + """Insight about agent expertise development""" + + agent_id: str + skill_area: str + insight_type: str # 'improving', 'declining', 'plateau', 'breakthrough' + description: str + confidence: float + evidence: Dict[str, Any] + timestamp: datetime + + +@dataclass +class AgentPerformanceMetrics: + """Performance metrics for an agent""" + + agent_id: str + total_tasks: int + success_rate: float + avg_expertise_level: float + improving_skills_count: int + declining_skills_count: int + memory_usage_stats: Dict[str, Any] + recent_performance_trend: str + top_skill_areas: List[Dict[str, Any]] + + +class AgentExpertiseTracker: + """ + Tracks and analyzes agent expertise development, providing insights + into learning patterns, performance trends, and optimization opportunities. + """ + + def __init__(self, database_url: str): + self.database_url = database_url + self.insights_cache: Dict[str, List[ExpertiseInsight]] = {} + self.metrics_cache: Dict[str, AgentPerformanceMetrics] = {} + self.cache_ttl = 300 # 5 minutes + self.last_cache_update = {} + + async def get_agent_performance_metrics( + self, agent_id: str + ) -> Optional[AgentPerformanceMetrics]: + """Get comprehensive performance metrics for an agent""" + + # Check cache first + if ( + agent_id in self.metrics_cache + and agent_id in self.last_cache_update + and (datetime.now() - self.last_cache_update[agent_id]).total_seconds() + < self.cache_ttl + ): + return self.metrics_cache[agent_id] + + try: + async with get_db_connection() as conn: + # Get basic performance stats + basic_stats = await conn.fetchrow( + """ + SELECT + COUNT(DISTINCT am.task_id) as total_tasks, + AVG(CASE WHEN am.memory_type = 'success' THEN 1.0 ELSE 0.0 END) as success_rate, + COUNT(DISTINCT am.id) as total_memories + FROM agent_memory am + WHERE am.agent_id = $1 + """, + agent_id, + ) + + # Get expertise summary + expertise_stats = await conn.fetchrow( + """ + SELECT + AVG(expertise_level) as avg_expertise_level, + COUNT(CASE WHEN performance_trend = 'improving' THEN 1 END) as improving_skills, + COUNT(CASE WHEN performance_trend = 'declining' THEN 1 END) as declining_skills + FROM agent_expertise + WHERE agent_id = $1 + """, + agent_id, + ) + + # Get memory usage statistics + memory_stats = await conn.fetchrow( + """ + SELECT + COUNT(*) as total_memories, + AVG(confidence_score) as avg_confidence, + SUM(usage_count) as total_usage, + COUNT(DISTINCT memory_type) as memory_types_used + FROM agent_memory + WHERE agent_id = $1 + """, + agent_id, + ) + + # Get top skill areas + top_skills = await conn.fetch( + """ + SELECT skill_area, expertise_level, success_rate, task_count, performance_trend + FROM agent_expertise + WHERE agent_id = $1 + ORDER BY expertise_level DESC, success_rate DESC + LIMIT 5 + """, + agent_id, + ) + + # Determine recent performance trend + recent_trend = await self._calculate_recent_trend(agent_id, conn) + + # Build metrics object + metrics = AgentPerformanceMetrics( + agent_id=agent_id, + total_tasks=basic_stats["total_tasks"] or 0, + success_rate=basic_stats["success_rate"] or 0.0, + avg_expertise_level=expertise_stats["avg_expertise_level"] or 0.0, + improving_skills_count=expertise_stats["improving_skills"] or 0, + declining_skills_count=expertise_stats["declining_skills"] or 0, + memory_usage_stats={ + "total_memories": memory_stats["total_memories"] or 0, + "avg_confidence": ( + float(memory_stats["avg_confidence"]) + if memory_stats["avg_confidence"] + else 0.0 + ), + "total_usage": memory_stats["total_usage"] or 0, + "memory_types_used": memory_stats["memory_types_used"] or 0, + }, + recent_performance_trend=recent_trend, + top_skill_areas=[dict(skill) for skill in top_skills], + ) + + # Cache the result + self.metrics_cache[agent_id] = metrics + self.last_cache_update[agent_id] = datetime.now() + + return metrics + + except Exception as e: + logger.error(f"Error getting performance metrics for agent {agent_id}: {e}") + return None + + async def generate_expertise_insights( + self, agent_id: str + ) -> List[ExpertiseInsight]: + """Generate insights about agent expertise development""" + + # Check cache first + if ( + agent_id in self.insights_cache + and agent_id in self.last_cache_update + and (datetime.now() - self.last_cache_update[agent_id]).total_seconds() + < self.cache_ttl + ): + return self.insights_cache[agent_id] + + insights = [] + + try: + async with get_db_connection() as conn: + # Analyze learning velocity patterns + learning_insights = await self._analyze_learning_velocity( + agent_id, conn + ) + insights.extend(learning_insights) + + # Analyze skill development patterns + skill_insights = await self._analyze_skill_development(agent_id, conn) + insights.extend(skill_insights) + + # Analyze memory usage patterns + memory_insights = await self._analyze_memory_patterns(agent_id, conn) + insights.extend(memory_insights) + + # Cache the results + self.insights_cache[agent_id] = insights + self.last_cache_update[agent_id] = datetime.now() + + except Exception as e: + logger.error(f"Error generating insights for agent {agent_id}: {e}") + + return insights + + async def get_system_wide_expertise_summary(self) -> Dict[str, Any]: + """Get system-wide expertise and performance summary""" + + try: + async with get_db_connection() as conn: + # Overall system stats + system_stats = await conn.fetchrow(""" + SELECT + COUNT(DISTINCT a.id) as total_agents, + COUNT(DISTINCT ae.skill_area) as total_skill_areas, + AVG(ae.expertise_level) as avg_system_expertise, + COUNT(CASE WHEN ae.performance_trend = 'improving' THEN 1 END) as improving_agents, + COUNT(CASE WHEN ae.performance_trend = 'declining' THEN 1 END) as declining_agents + FROM agents a + LEFT JOIN agent_expertise ae ON a.id = ae.agent_id + """) + + # Memory system stats + memory_stats = await conn.fetchrow(""" + SELECT + COUNT(*) as total_memories, + AVG(confidence_score) as avg_confidence, + SUM(usage_count) as total_usage, + COUNT(DISTINCT agent_id) as agents_with_memory + FROM agent_memory + """) + + # Top performing skill areas + top_skill_areas = await conn.fetch(""" + SELECT + skill_area, + COUNT(*) as agent_count, + AVG(expertise_level) as avg_expertise, + AVG(success_rate) as avg_success_rate + FROM agent_expertise + GROUP BY skill_area + ORDER BY avg_expertise DESC, avg_success_rate DESC + LIMIT 10 + """) + + # Recent activity + recent_activity = await conn.fetchrow(""" + SELECT + COUNT(CASE WHEN created_at > NOW() - INTERVAL '24 hours' THEN 1 END) as memories_24h, + COUNT(CASE WHEN created_at > NOW() - INTERVAL '7 days' THEN 1 END) as memories_7d, + COUNT(DISTINCT CASE WHEN created_at > NOW() - INTERVAL '24 hours' THEN agent_id END) as active_agents_24h + FROM agent_memory + """) + + return { + "system_stats": dict(system_stats) if system_stats else {}, + "memory_stats": dict(memory_stats) if memory_stats else {}, + "top_skill_areas": [dict(skill) for skill in top_skill_areas], + "recent_activity": dict(recent_activity) if recent_activity else {}, + "timestamp": datetime.now().isoformat(), + } + + except Exception as e: + logger.error(f"Error getting system-wide expertise summary: {e}") + return {"error": str(e)} + + async def _calculate_recent_trend(self, agent_id: str, conn) -> str: + """Calculate recent performance trend for an agent""" + + try: + # Get recent task outcomes + recent_outcomes = await conn.fetch( + """ + SELECT + DATE_TRUNC('day', created_at) as date, + AVG(CASE WHEN memory_type = 'success' THEN 1.0 ELSE 0.0 END) as daily_success_rate + FROM agent_memory + WHERE agent_id = $1 + AND created_at > NOW() - INTERVAL '14 days' + AND memory_type IN ('success', 'task_outcome') + GROUP BY DATE_TRUNC('day', created_at) + ORDER BY date DESC + LIMIT 7 + """, + agent_id, + ) + + if len(recent_outcomes) < 3: + return "insufficient_data" + + # Calculate trend + success_rates = [ + float(row["daily_success_rate"]) for row in recent_outcomes + ] + + # Simple linear trend calculation + if len(success_rates) >= 3: + early_avg = sum(success_rates[-3:]) / 3 + recent_avg = sum(success_rates[:3]) / 3 + + if recent_avg > early_avg + 0.1: + return "improving" + elif recent_avg < early_avg - 0.1: + return "declining" + else: + return "stable" + + return "stable" + + except Exception as e: + logger.error(f"Error calculating recent trend: {e}") + return "unknown" + + async def _analyze_learning_velocity( + self, agent_id: str, conn + ) -> List[ExpertiseInsight]: + """Analyze learning velocity patterns""" + + insights = [] + + try: + # Get skills with high learning velocity + fast_learners = await conn.fetch( + """ + SELECT skill_area, learning_velocity, expertise_level, task_count + FROM agent_expertise + WHERE agent_id = $1 AND learning_velocity > 0.1 + ORDER BY learning_velocity DESC + """, + agent_id, + ) + + for skill in fast_learners: + insights.append( + ExpertiseInsight( + agent_id=agent_id, + skill_area=skill["skill_area"], + insight_type="improving", + description=f"Rapid improvement in {skill['skill_area']} with velocity {skill['learning_velocity']:.2f}", + confidence=0.8, + evidence={ + "learning_velocity": float(skill["learning_velocity"]), + "expertise_level": float(skill["expertise_level"]), + "task_count": skill["task_count"], + }, + timestamp=datetime.now(), + ) + ) + + # Get skills with declining performance + declining_skills = await conn.fetch( + """ + SELECT skill_area, learning_velocity, expertise_level, task_count + FROM agent_expertise + WHERE agent_id = $1 AND learning_velocity < -0.05 + ORDER BY learning_velocity ASC + """, + agent_id, + ) + + for skill in declining_skills: + insights.append( + ExpertiseInsight( + agent_id=agent_id, + skill_area=skill["skill_area"], + insight_type="declining", + description=f"Performance decline in {skill['skill_area']} - may need attention", + confidence=0.7, + evidence={ + "learning_velocity": float(skill["learning_velocity"]), + "expertise_level": float(skill["expertise_level"]), + "task_count": skill["task_count"], + }, + timestamp=datetime.now(), + ) + ) + + except Exception as e: + logger.error(f"Error analyzing learning velocity: {e}") + + return insights + + async def _analyze_skill_development( + self, agent_id: str, conn + ) -> List[ExpertiseInsight]: + """Analyze skill development patterns""" + + insights = [] + + try: + # Find breakthrough moments (significant expertise jumps) + breakthroughs = await conn.fetch( + """ + SELECT skill_area, expertise_level, success_rate, task_count + FROM agent_expertise + WHERE agent_id = $1 + AND expertise_level > 0.7 + AND success_rate > 0.8 + AND task_count >= 5 + """, + agent_id, + ) + + for breakthrough in breakthroughs: + insights.append( + ExpertiseInsight( + agent_id=agent_id, + skill_area=breakthrough["skill_area"], + insight_type="breakthrough", + description=f"Expert level achieved in {breakthrough['skill_area']} with {breakthrough['success_rate']:.1%} success rate", + confidence=0.9, + evidence={ + "expertise_level": float(breakthrough["expertise_level"]), + "success_rate": float(breakthrough["success_rate"]), + "task_count": breakthrough["task_count"], + }, + timestamp=datetime.now(), + ) + ) + + # Find plateau situations (high task count but low expertise) + plateaus = await conn.fetch( + """ + SELECT skill_area, expertise_level, success_rate, task_count + FROM agent_expertise + WHERE agent_id = $1 + AND task_count > 10 + AND expertise_level < 0.4 + AND learning_velocity BETWEEN -0.02 AND 0.02 + """, + agent_id, + ) + + for plateau in plateaus: + insights.append( + ExpertiseInsight( + agent_id=agent_id, + skill_area=plateau["skill_area"], + insight_type="plateau", + description=f"Learning plateau in {plateau['skill_area']} - consider new approaches", + confidence=0.6, + evidence={ + "expertise_level": float(plateau["expertise_level"]), + "success_rate": float(plateau["success_rate"]), + "task_count": plateau["task_count"], + }, + timestamp=datetime.now(), + ) + ) + + except Exception as e: + logger.error(f"Error analyzing skill development: {e}") + + return insights + + async def _analyze_memory_patterns( + self, agent_id: str, conn + ) -> List[ExpertiseInsight]: + """Analyze memory usage and effectiveness patterns""" + + insights = [] + + try: + # Analyze memory types and their effectiveness + memory_effectiveness = await conn.fetchrow( + """ + SELECT + COUNT(*) as total_memories, + AVG(usage_count) as avg_usage, + AVG(confidence_score) as avg_confidence, + COUNT(CASE WHEN usage_count > 5 THEN 1 END) as high_usage_memories + FROM agent_memory + WHERE agent_id = $1 + """, + agent_id, + ) + + if memory_effectiveness and memory_effectiveness["total_memories"] > 50: + high_usage_ratio = ( + memory_effectiveness["high_usage_memories"] + / memory_effectiveness["total_memories"] + ) + + if ( + high_usage_ratio > 0.2 + ): # More than 20% of memories are highly reused + insights.append( + ExpertiseInsight( + agent_id=agent_id, + skill_area="memory_management", + insight_type="improving", + description=f"Excellent memory reuse patterns - {high_usage_ratio:.1%} of memories are frequently accessed", + confidence=0.8, + evidence={ + "total_memories": memory_effectiveness[ + "total_memories" + ], + "high_usage_ratio": high_usage_ratio, + "avg_confidence": float( + memory_effectiveness["avg_confidence"] + ), + }, + timestamp=datetime.now(), + ) + ) + + except Exception as e: + logger.error(f"Error analyzing memory patterns: {e}") + + return insights + + async def clear_cache(self, agent_id: Optional[str] = None): + """Clear analytics cache""" + if agent_id: + self.insights_cache.pop(agent_id, None) + self.metrics_cache.pop(agent_id, None) + self.last_cache_update.pop(agent_id, None) + else: + self.insights_cache.clear() + self.metrics_cache.clear() + self.last_cache_update.clear() diff --git a/services/orchestrator/claude_code_wrapper.py b/services/orchestrator/claude_code_wrapper.py index f1186e5..5188ab8 100644 --- a/services/orchestrator/claude_code_wrapper.py +++ b/services/orchestrator/claude_code_wrapper.py @@ -1,842 +1,842 @@ -import asyncio -import json -import os -import subprocess # nosec B404 -- used only with static executable names + arg lists and shell=False (see _run_tests) -import tempfile -import time -from pathlib import Path -from typing import Any, Dict, List, Optional, Type - -# Import Anthropic SDK for real Claude integration -import anthropic -from anthropic import Anthropic -from crewai.tools import BaseTool -from pydantic import BaseModel, Field - -# Import conversation manager for full chat tracking. -# This module is imported both as part of the `services.orchestrator` package -# (relative form, e.g. from main.py/agent_manager.py) and flat with -# services/orchestrator on sys.path (e.g. from tests and main_with_hierarchy.py), -# so support both — mirrors the existing pattern in hierarchy_endpoints.py. -try: - from .conversation_manager import ConversationManager, MessageType -except ImportError: # pragma: no cover - flat import (no parent package) - from conversation_manager import ConversationManager, MessageType - - -class ClaudeCodeInput(BaseModel): - """Input schema for Claude Code tool""" - - task: str = Field(description="Coding task to complete") - language: str = Field(default="python", description="Programming language") - context: str = Field(default="", description="Additional context or requirements") - include_tests: bool = Field( - default=True, description="Whether to include unit tests" - ) - include_docs: bool = Field( - default=True, description="Whether to include documentation" - ) - file_path: Optional[str] = Field( - default=None, description="Optional file path for code context" - ) - - -class ClaudeCodeWrapper(BaseTool): - name: str = "claude_code" - description: str = """ - Execute advanced coding tasks using Claude AI with real-time code generation, - testing, and documentation. Supports multiple programming languages and - follows industry best practices. Enhanced for repository context and Git integration. - """ - args_schema: Type[BaseModel] = ClaudeCodeInput - - # Runtime attributes. ``BaseTool`` is a Pydantic v2 model, which rejects - # assignment to undeclared attributes ("object has no field ..."). These - # are declared as model fields (rather than PrivateAttr) so they remain - # publicly readable on the instance (e.g. ``wrapper.client`` / - # ``wrapper.model``), preserving the tool's public interface. Object - # handles (Anthropic SDK client, git/conversation managers) are typed - # ``Any`` so Pydantic stores them as-is without schema validation. - client: Optional[Any] = None - model: str = "claude-3-5-sonnet-20241022" - workspace_path: Optional[str] = None - git_manager: Optional[Any] = None - agent_id: Optional[str] = None - task_id: Optional[str] = None - conversation_manager: Optional[Any] = None - conversation_session_id: Optional[str] = None - current_context: Dict[str, Any] = Field(default_factory=dict) - repository_context: Dict[str, Any] = Field(default_factory=dict) - - def __init__( - self, - workspace_path: Optional[str] = None, - git_manager: Optional[Any] = None, - agent_id: Optional[str] = None, - task_id: Optional[str] = None, - conversation_manager: Optional[ConversationManager] = None, - ): - super().__init__() - self.client = Anthropic( - api_key=os.getenv("ANTHROPIC_API_KEY"), - ) - self.model = "claude-3-5-sonnet-20241022" - self.workspace_path = workspace_path or os.getcwd() - self.git_manager = git_manager - self.agent_id = agent_id - self.task_id = task_id - self.conversation_manager = conversation_manager or ConversationManager() - self.conversation_session_id: Optional[str] = None - self.current_context = {} # Store context between iterations - - # Repository context - self.repository_context = { - "files_changed": [], - "current_branch": None, - "last_commit": None, - "iteration_count": 0, - } - - def _run( - self, - task: str, - language: str = "python", - context: str = "", - include_tests: bool = True, - include_docs: bool = True, - file_path: Optional[str] = None, - iteration_number: Optional[int] = None, - ) -> str: - """Execute Claude Code for a specific task with real AI integration""" - - try: - # Update iteration count - if iteration_number: - self.repository_context["iteration_count"] = iteration_number - else: - self.repository_context["iteration_count"] += 1 - - # Get repository context if Git manager is available - repo_context = "" - if self.git_manager: - try: - # Note: In a full implementation, we'd make this method async - # For now, we'll skip the Git context in the sync version - repo_context = ( - "Repository context: Available (Git manager configured)" - ) - except Exception as e: - repo_context = f"Repository context unavailable: {str(e)}" - - # Read existing file context if provided - existing_code = "" - if file_path: - # Use workspace-relative path if available - full_path = ( - os.path.join(self.workspace_path, file_path) - if not os.path.isabs(file_path) - else file_path - ) - if os.path.exists(full_path): - with open(full_path, "r") as f: - existing_code = f.read() - - # Prepare the comprehensive prompt with repository context - prompt = self._build_prompt( - task=task, - language=language, - context=context, - existing_code=existing_code, - include_tests=include_tests, - include_docs=include_docs, - repo_context=repo_context, - ) - - # Call Claude API - response = self.client.messages.create( - model=self.model, - max_tokens=4096, - temperature=0.3, # Lower temperature for more consistent code - messages=[{"role": "user", "content": prompt}], - ) - - # Parse the response and extract code files - result = self._parse_response( - response.content[0].text, language, include_tests, include_docs - ) - - # Save files to workspace if available, otherwise use temp directory - if self.workspace_path and os.path.exists(self.workspace_path): - saved_files = self._save_files_to_workspace(result["files"]) - else: - with tempfile.TemporaryDirectory() as tmpdir: - saved_files = self._save_files(result["files"], tmpdir) - - # Update repository context with changed files - self.repository_context["files_changed"].extend( - [ - f["filename"] - for f in result["files"] - if f["type"] == "implementation" - ] - ) - - # Run tests if generated and in workspace - test_results = None - if include_tests and any(f["type"] == "test" for f in result["files"]): - if self.workspace_path and os.path.exists(self.workspace_path): - test_results = self._run_tests(self.workspace_path, language) - else: - with tempfile.TemporaryDirectory() as tmpdir: - self._save_files(result["files"], tmpdir) - test_results = self._run_tests(tmpdir, language) - - return json.dumps( - { - "status": "success", - "files": result["files"], - "explanation": result.get("explanation", ""), - "test_results": test_results, - "commit_message": result.get("commit_message", ""), - "execution_summary": f"Generated {len(result['files'])} files for {language} task: {task[:100]}...", - "iteration": self.repository_context["iteration_count"], - "workspace_path": self.workspace_path, - "repository_context": self.repository_context, - } - ) - - except anthropic.APIError as e: - return json.dumps( - { - "status": "error", - "error": f"Claude API error: {str(e)}", - "error_type": "api_error", - } - ) - except Exception as e: - return json.dumps( - { - "status": "error", - "error": f"Unexpected error: {str(e)}", - "error_type": "general_error", - } - ) - - def _build_prompt( - self, - task: str, - language: str, - context: str, - existing_code: str, - include_tests: bool, - include_docs: bool, - repo_context: str = "", - ) -> str: - """Build a comprehensive prompt for Claude with repository context""" - - # Build agent context - agent_info = "" - if self.agent_id and self.task_id: - agent_info = f""" -**Agent Context**: -- Agent ID: {self.agent_id} -- Task ID: {self.task_id} -- Iteration: {self.repository_context['iteration_count']} -- Workspace: {self.workspace_path} -""" - - prompt = f""" -You are an expert {language} developer working autonomously as part of FuzeAgent AI team. I need you to complete the following coding task: - -{agent_info} - -**Task**: {task} - -**Programming Language**: {language} - -**Additional Context**: {context} - -{repo_context} - -**Existing Code** (if any): -```{language} -{existing_code} -``` - -**Requirements**: -1. Write clean, maintainable, and well-documented code -2. Follow {language} best practices and conventions -3. Include proper error handling -4. Use type hints (where applicable) -5. {"Include comprehensive unit tests" if include_tests else "Focus only on implementation"} -6. {"Include docstrings and comments" if include_docs else "Minimal documentation"} -7. Consider the repository context and maintain consistency with existing code -8. Write code that integrates well with the current branch and recent changes - -**Output Format**: -Please structure your response as follows: - -## Explanation -Brief explanation of your approach and key decisions, considering the repository context. - -## Implementation - -### Main Code -```{language} -# Your main implementation here -``` - -{"### Tests" if include_tests else ""} -{f"```{language}" if include_tests else ""} -{"# Your test code here" if include_tests else ""} -{f"```" if include_tests else ""} - -{"### Documentation" if include_docs else ""} -{"```markdown" if include_docs else ""} -{"# Your documentation here" if include_docs else ""} -{f"```" if include_docs else ""} - -## Commit Message -Suggest a concise git commit message for these changes that follows the repository's commit history style. - -Please ensure the code is production-ready and follows industry standards. -""" - return prompt - - def _parse_response( - self, response: str, language: str, include_tests: bool, include_docs: bool - ) -> Dict[str, Any]: - """Parse Claude's response and extract code files""" - - files = [] - explanation = "" - commit_message = "" - - # Extract explanation - if "## Explanation" in response: - explanation_start = response.find("## Explanation") + len("## Explanation") - explanation_end = response.find("## Implementation") - if explanation_end > explanation_start: - explanation = response[explanation_start:explanation_end].strip() - - # Extract commit message - if "## Commit Message" in response: - commit_start = response.find("## Commit Message") + len("## Commit Message") - commit_message = response[commit_start:].strip() - # Clean up the commit message - commit_message = commit_message.split("\n")[0].strip() - - # Extract main code - main_code = self._extract_code_block(response, "### Main Code", language) - if main_code: - file_ext = self._get_file_extension(language) - files.append( - { - "filename": f"main.{file_ext}", - "content": main_code, - "type": "implementation", - "language": language, - } - ) - - # Extract tests if requested - if include_tests: - test_code = self._extract_code_block(response, "### Tests", language) - if test_code: - test_ext = self._get_file_extension(language) - files.append( - { - "filename": f"test_main.{test_ext}", - "content": test_code, - "type": "test", - "language": language, - } - ) - - # Extract documentation if requested - if include_docs: - docs = self._extract_code_block(response, "### Documentation", "markdown") - if docs: - files.append( - { - "filename": "README.md", - "content": docs, - "type": "documentation", - "language": "markdown", - } - ) - - return { - "files": files, - "explanation": explanation, - "commit_message": commit_message, - } - - def _extract_code_block( - self, text: str, section: str, language: str - ) -> Optional[str]: - """Extract code block from a specific section""" - - section_start = text.find(section) - if section_start == -1: - return None - - # Find the start of the code block - code_start = text.find(f"```{language}", section_start) - if code_start == -1: - code_start = text.find("```", section_start) - if code_start == -1: - return None - - # Find the end of the code block - code_content_start = text.find("\n", code_start) + 1 - code_end = text.find("```", code_content_start) - - if code_end == -1: - return None - - return text[code_content_start:code_end].strip() - - def _get_file_extension(self, language: str) -> str: - """Get appropriate file extension for language""" - extensions = { - "python": "py", - "javascript": "js", - "typescript": "ts", - "java": "java", - "cpp": "cpp", - "c": "c", - "rust": "rs", - "go": "go", - "ruby": "rb", - "php": "php", - "swift": "swift", - "kotlin": "kt", - "scala": "scala", - "r": "R", - "sql": "sql", - "html": "html", - "css": "css", - "shell": "sh", - "bash": "sh", - } - return extensions.get(language.lower(), "txt") - - def _save_files(self, files: List[Dict], tmpdir: str) -> List[str]: - """Save generated files to temporary directory""" - saved_files = [] - - for file_info in files: - file_path = os.path.join(tmpdir, file_info["filename"]) - with open(file_path, "w") as f: - f.write(file_info["content"]) - saved_files.append(file_path) - - return saved_files - - def _run_tests(self, tmpdir: str, language: str) -> Optional[Dict[str, Any]]: - """Run tests for the generated code""" - - try: - if language == "python": - # Try to run pytest - result = subprocess.run( # nosec B603 B607 -- shell=False, static argv; runs pytest on generated code inside an isolated tmpdir - ["python", "-m", "pytest", tmpdir, "-v"], - capture_output=True, - text=True, - timeout=60, - cwd=tmpdir, - ) - - return { - "exit_code": result.returncode, - "stdout": result.stdout, - "stderr": result.stderr, - "success": result.returncode == 0, - } - elif language == "javascript": - # Try to run with node - test_files = [f for f in os.listdir(tmpdir) if f.startswith("test_")] - if test_files: - result = subprocess.run( # nosec B603 B607 -- shell=False, static argv; runs a generated test file inside an isolated tmpdir - ["node", test_files[0]], - capture_output=True, - text=True, - timeout=60, - cwd=tmpdir, - ) - - return { - "exit_code": result.returncode, - "stdout": result.stdout, - "stderr": result.stderr, - "success": result.returncode == 0, - } - - return None - - except subprocess.TimeoutExpired: - return { - "exit_code": -1, - "stdout": "", - "stderr": "Test execution timed out", - "success": False, - } - except Exception as e: - return { - "exit_code": -1, - "stdout": "", - "stderr": f"Test execution error: {str(e)}", - "success": False, - } - - def _build_repository_context( - self, branch_status: Dict[str, Any], commit_history: List[Any] - ) -> str: - """Build repository context string for the prompt""" - context_parts = ["**Repository Context**:"] - - if branch_status: - current_branch = branch_status.get("current_branch", "unknown") - feature_branch = branch_status.get("feature_branch") - has_changes = branch_status.get("has_uncommitted_changes", False) - remote_status = branch_status.get("remote_status", "unknown") - - context_parts.append(f"- Current Branch: `{current_branch}`") - if feature_branch: - context_parts.append(f"- Feature Branch: `{feature_branch}`") - context_parts.append( - f"- Uncommitted Changes: {'Yes' if has_changes else 'No'}" - ) - context_parts.append(f"- Remote Status: {remote_status}") - - if commit_history: - context_parts.append("- Recent Commits:") - for i, commit in enumerate(commit_history[:3]): - context_parts.append(f" {i+1}. `{commit.hash[:8]}` - {commit.message}") - if commit.files_changed: - context_parts.append( - f" Files: {', '.join(commit.files_changed[:5])}" - ) - - if self.repository_context.get("files_changed"): - changed_files = list(set(self.repository_context["files_changed"])) - context_parts.append( - f"- Files Modified This Session: {', '.join(changed_files)}" - ) - - return "\n".join(context_parts) + "\n" - - def _save_files_to_workspace(self, files: List[Dict]) -> List[str]: - """Save generated files directly to workspace""" - saved_files = [] - - for file_info in files: - file_path = os.path.join(self.workspace_path, file_info["filename"]) - - # Create directory if needed - os.makedirs(os.path.dirname(file_path), exist_ok=True) - - with open(file_path, "w") as f: - f.write(file_info["content"]) - saved_files.append(file_path) - - return saved_files - - async def commit_and_push_changes( - self, commit_message: str, files: Optional[List[str]] = None - ) -> Dict[str, Any]: - """Commit and push changes using Git manager""" - if not self.git_manager: - return {"success": False, "error": "No Git manager available"} - - try: - # Commit changes - commit_hash = await self.git_manager.commit_changes( - message=commit_message, - files=files, - iteration_number=self.repository_context["iteration_count"], - ) - - if commit_hash: - self.repository_context["last_commit"] = commit_message - return { - "success": True, - "commit_hash": commit_hash, - "message": "Changes committed successfully", - } - else: - return {"success": True, "message": "No changes to commit"} - - except Exception as e: - return {"success": False, "error": f"Failed to commit changes: {str(e)}"} - - def get_repository_context(self) -> Dict[str, Any]: - """Get current repository context""" - return self.repository_context.copy() - - def reset_context(self): - """Reset the repository context""" - self.repository_context = { - "files_changed": [], - "current_branch": None, - "last_commit": None, - "iteration_count": 0, - } - - async def start_conversation_session(self, sandbox_id: str) -> str: - """Start a conversation session for tracking all Claude Code interactions""" - if not self.agent_id or not self.task_id: - raise ValueError("Agent ID and Task ID required for conversation tracking") - - self.conversation_session_id = ( - await self.conversation_manager.start_conversation_session( - agent_id=self.agent_id, - task_id=self.task_id, - sandbox_id=sandbox_id, - metadata={ - "workspace_path": self.workspace_path, - "model": self.model, - "git_enabled": bool(self.git_manager), - }, - ) - ) - return self.conversation_session_id - - async def end_conversation_session(self) -> bool: - """End the current conversation session""" - if not self.conversation_session_id: - return False - - success = await self.conversation_manager.end_conversation_session( - self.conversation_session_id - ) - self.conversation_session_id = None - return success - - async def execute_task_async( - self, - task: str, - language: str = "python", - context: str = "", - include_tests: bool = True, - include_docs: bool = True, - file_path: Optional[str] = None, - iteration_number: Optional[int] = None, - ) -> Dict[str, Any]: - """Async version of task execution with full conversation tracking""" - - try: - # Update iteration count - if iteration_number: - self.repository_context["iteration_count"] = iteration_number - else: - self.repository_context["iteration_count"] += 1 - - current_iteration = self.repository_context["iteration_count"] - - # Get repository context if Git manager is available - repo_context = "" - if self.git_manager: - try: - branch_status = await self.git_manager.get_branch_status() - self.repository_context["current_branch"] = branch_status.get( - "current_branch" - ) - - commit_history = await self.git_manager.get_commit_history(limit=3) - if commit_history: - self.repository_context["last_commit"] = commit_history[ - 0 - ].message - - repo_context = self._build_repository_context( - branch_status, commit_history - ) - except Exception as e: - repo_context = f"Repository context unavailable: {str(e)}" - - # Read existing file context if provided - existing_code = "" - if file_path: - # Use workspace-relative path if available - full_path = ( - os.path.join(self.workspace_path, file_path) - if not os.path.isabs(file_path) - else file_path - ) - if os.path.exists(full_path): - with open(full_path, "r") as f: - existing_code = f.read() - - # Prepare the comprehensive prompt with repository context - prompt = self._build_prompt( - task=task, - language=language, - context=context, - existing_code=existing_code, - include_tests=include_tests, - include_docs=include_docs, - repo_context=repo_context, - ) - - # Store user prompt in conversation history - if self.conversation_session_id and self.task_id: - await self.conversation_manager.store_user_prompt( - session_id=self.conversation_session_id, - task_id=self.task_id, - iteration_number=current_iteration, - prompt=prompt, - model=self.model, - temperature=0.3, - metadata={ - "task_description": ( - task[:200] + "..." if len(task) > 200 else task - ), - "language": language, - "include_tests": include_tests, - "include_docs": include_docs, - "file_path": file_path, - "workspace_path": self.workspace_path, - }, - ) - - # Record start time for response time tracking - start_time = time.time() - - # Call Claude API - response = self.client.messages.create( - model=self.model, - max_tokens=4096, - temperature=0.3, # Lower temperature for more consistent code - messages=[{"role": "user", "content": prompt}], - ) - - # Extract response content and token usage - response_content = response.content[0].text - token_count = ( - getattr(response.usage, "output_tokens", None) - if hasattr(response, "usage") - else None - ) - - # Store Claude response in conversation history - if self.conversation_session_id and self.task_id: - await self.conversation_manager.store_claude_response( - session_id=self.conversation_session_id, - task_id=self.task_id, - iteration_number=current_iteration, - response=response_content, - token_count=token_count, - model=self.model, - start_time=start_time, - metadata={ - "prompt_length": len(prompt), - "response_length": len(response_content), - }, - ) - - # Parse the response and extract code files - result = self._parse_response( - response_content, language, include_tests, include_docs - ) - - # Save files to workspace if available - saved_files = [] - if self.workspace_path and os.path.exists(self.workspace_path): - saved_files = self._save_files_to_workspace(result["files"]) - - # Store code generations in database - if self.task_id: - for file_info in result["files"]: - await self.conversation_manager.store_code_generation( - task_id=self.task_id, - iteration_number=current_iteration, - file_path=file_info["filename"], - file_type=file_info["type"], - language=file_info.get("language", language), - content=file_info["content"], - ) - - # Update repository context with changed files - self.repository_context["files_changed"].extend( - [ - f["filename"] - for f in result["files"] - if f["type"] == "implementation" - ] - ) - - # Run tests if generated and in workspace - test_results = None - if include_tests and any(f["type"] == "test" for f in result["files"]): - if self.workspace_path and os.path.exists(self.workspace_path): - test_results = self._run_tests(self.workspace_path, language) - - # Store test results - if self.conversation_session_id and self.task_id: - await self.conversation_manager.store_message( - session_id=self.conversation_session_id, - message={ - "task_id": self.task_id, - "iteration_number": current_iteration, - "message_type": MessageType.TEST_RESULT, - "content": json.dumps(test_results), - "metadata": { - "test_framework": ( - "pytest" if language == "python" else "jest" - ), - "workspace_path": self.workspace_path, - }, - }, - ) - - return { - "status": "success", - "files": result["files"], - "saved_files": saved_files, - "explanation": result.get("explanation", ""), - "test_results": test_results, - "commit_message": result.get("commit_message", ""), - "execution_summary": f"Generated {len(result['files'])} files for {language} task: {task[:100]}...", - "iteration": current_iteration, - "workspace_path": self.workspace_path, - "repository_context": self.repository_context, - "conversation_tracked": bool(self.conversation_session_id), - "token_count": token_count, - } - - except Exception as e: - # Store error in conversation history - if self.conversation_session_id and self.task_id: - try: - await self.conversation_manager.store_message( - session_id=self.conversation_session_id, - message={ - "task_id": self.task_id, - "iteration_number": self.repository_context[ - "iteration_count" - ], - "message_type": MessageType.ERROR_MESSAGE, - "content": str(e), - "metadata": { - "error_type": type(e).__name__, - "task_description": ( - task[:200] + "..." if len(task) > 200 else task - ), - }, - }, - ) - except Exception as conv_error: - # Don't let conversation storage errors break the main flow - print( - f"Warning: Failed to store error in conversation: {conv_error}" - ) - - return { - "status": "error", - "error": str(e), - "error_type": "execution_error", - "iteration": self.repository_context["iteration_count"], - "conversation_tracked": bool(self.conversation_session_id), - } +import asyncio +import json +import os +import subprocess # nosec B404 -- used only with static executable names + arg lists and shell=False (see _run_tests) +import tempfile +import time +from pathlib import Path +from typing import Any, Dict, List, Optional, Type + +# Import Anthropic SDK for real Claude integration +import anthropic +from anthropic import Anthropic +from crewai.tools import BaseTool +from pydantic import BaseModel, Field + +# Import conversation manager for full chat tracking. +# This module is imported both as part of the `services.orchestrator` package +# (relative form, e.g. from main.py/agent_manager.py) and flat with +# services/orchestrator on sys.path (e.g. from tests and main_with_hierarchy.py), +# so support both — mirrors the existing pattern in hierarchy_endpoints.py. +try: + from .conversation_manager import ConversationManager, MessageType +except ImportError: # pragma: no cover - flat import (no parent package) + from conversation_manager import ConversationManager, MessageType + + +class ClaudeCodeInput(BaseModel): + """Input schema for Claude Code tool""" + + task: str = Field(description="Coding task to complete") + language: str = Field(default="python", description="Programming language") + context: str = Field(default="", description="Additional context or requirements") + include_tests: bool = Field( + default=True, description="Whether to include unit tests" + ) + include_docs: bool = Field( + default=True, description="Whether to include documentation" + ) + file_path: Optional[str] = Field( + default=None, description="Optional file path for code context" + ) + + +class ClaudeCodeWrapper(BaseTool): + name: str = "claude_code" + description: str = """ + Execute advanced coding tasks using Claude AI with real-time code generation, + testing, and documentation. Supports multiple programming languages and + follows industry best practices. Enhanced for repository context and Git integration. + """ + args_schema: Type[BaseModel] = ClaudeCodeInput + + # Runtime attributes. ``BaseTool`` is a Pydantic v2 model, which rejects + # assignment to undeclared attributes ("object has no field ..."). These + # are declared as model fields (rather than PrivateAttr) so they remain + # publicly readable on the instance (e.g. ``wrapper.client`` / + # ``wrapper.model``), preserving the tool's public interface. Object + # handles (Anthropic SDK client, git/conversation managers) are typed + # ``Any`` so Pydantic stores them as-is without schema validation. + client: Optional[Any] = None + model: str = "claude-3-5-sonnet-20241022" + workspace_path: Optional[str] = None + git_manager: Optional[Any] = None + agent_id: Optional[str] = None + task_id: Optional[str] = None + conversation_manager: Optional[Any] = None + conversation_session_id: Optional[str] = None + current_context: Dict[str, Any] = Field(default_factory=dict) + repository_context: Dict[str, Any] = Field(default_factory=dict) + + def __init__( + self, + workspace_path: Optional[str] = None, + git_manager: Optional[Any] = None, + agent_id: Optional[str] = None, + task_id: Optional[str] = None, + conversation_manager: Optional[ConversationManager] = None, + ): + super().__init__() + self.client = Anthropic( + api_key=os.getenv("ANTHROPIC_API_KEY"), + ) + self.model = "claude-3-5-sonnet-20241022" + self.workspace_path = workspace_path or os.getcwd() + self.git_manager = git_manager + self.agent_id = agent_id + self.task_id = task_id + self.conversation_manager = conversation_manager or ConversationManager() + self.conversation_session_id: Optional[str] = None + self.current_context = {} # Store context between iterations + + # Repository context + self.repository_context = { + "files_changed": [], + "current_branch": None, + "last_commit": None, + "iteration_count": 0, + } + + def _run( + self, + task: str, + language: str = "python", + context: str = "", + include_tests: bool = True, + include_docs: bool = True, + file_path: Optional[str] = None, + iteration_number: Optional[int] = None, + ) -> str: + """Execute Claude Code for a specific task with real AI integration""" + + try: + # Update iteration count + if iteration_number: + self.repository_context["iteration_count"] = iteration_number + else: + self.repository_context["iteration_count"] += 1 + + # Get repository context if Git manager is available + repo_context = "" + if self.git_manager: + try: + # Note: In a full implementation, we'd make this method async + # For now, we'll skip the Git context in the sync version + repo_context = ( + "Repository context: Available (Git manager configured)" + ) + except Exception as e: + repo_context = f"Repository context unavailable: {str(e)}" + + # Read existing file context if provided + existing_code = "" + if file_path: + # Use workspace-relative path if available + full_path = ( + os.path.join(self.workspace_path, file_path) + if not os.path.isabs(file_path) + else file_path + ) + if os.path.exists(full_path): + with open(full_path, "r") as f: + existing_code = f.read() + + # Prepare the comprehensive prompt with repository context + prompt = self._build_prompt( + task=task, + language=language, + context=context, + existing_code=existing_code, + include_tests=include_tests, + include_docs=include_docs, + repo_context=repo_context, + ) + + # Call Claude API + response = self.client.messages.create( + model=self.model, + max_tokens=4096, + temperature=0.3, # Lower temperature for more consistent code + messages=[{"role": "user", "content": prompt}], + ) + + # Parse the response and extract code files + result = self._parse_response( + response.content[0].text, language, include_tests, include_docs + ) + + # Save files to workspace if available, otherwise use temp directory + if self.workspace_path and os.path.exists(self.workspace_path): + saved_files = self._save_files_to_workspace(result["files"]) + else: + with tempfile.TemporaryDirectory() as tmpdir: + saved_files = self._save_files(result["files"], tmpdir) + + # Update repository context with changed files + self.repository_context["files_changed"].extend( + [ + f["filename"] + for f in result["files"] + if f["type"] == "implementation" + ] + ) + + # Run tests if generated and in workspace + test_results = None + if include_tests and any(f["type"] == "test" for f in result["files"]): + if self.workspace_path and os.path.exists(self.workspace_path): + test_results = self._run_tests(self.workspace_path, language) + else: + with tempfile.TemporaryDirectory() as tmpdir: + self._save_files(result["files"], tmpdir) + test_results = self._run_tests(tmpdir, language) + + return json.dumps( + { + "status": "success", + "files": result["files"], + "explanation": result.get("explanation", ""), + "test_results": test_results, + "commit_message": result.get("commit_message", ""), + "execution_summary": f"Generated {len(result['files'])} files for {language} task: {task[:100]}...", + "iteration": self.repository_context["iteration_count"], + "workspace_path": self.workspace_path, + "repository_context": self.repository_context, + } + ) + + except anthropic.APIError as e: + return json.dumps( + { + "status": "error", + "error": f"Claude API error: {str(e)}", + "error_type": "api_error", + } + ) + except Exception as e: + return json.dumps( + { + "status": "error", + "error": f"Unexpected error: {str(e)}", + "error_type": "general_error", + } + ) + + def _build_prompt( + self, + task: str, + language: str, + context: str, + existing_code: str, + include_tests: bool, + include_docs: bool, + repo_context: str = "", + ) -> str: + """Build a comprehensive prompt for Claude with repository context""" + + # Build agent context + agent_info = "" + if self.agent_id and self.task_id: + agent_info = f""" +**Agent Context**: +- Agent ID: {self.agent_id} +- Task ID: {self.task_id} +- Iteration: {self.repository_context['iteration_count']} +- Workspace: {self.workspace_path} +""" + + prompt = f""" +You are an expert {language} developer working autonomously as part of FuzeAgent AI team. I need you to complete the following coding task: + +{agent_info} + +**Task**: {task} + +**Programming Language**: {language} + +**Additional Context**: {context} + +{repo_context} + +**Existing Code** (if any): +```{language} +{existing_code} +``` + +**Requirements**: +1. Write clean, maintainable, and well-documented code +2. Follow {language} best practices and conventions +3. Include proper error handling +4. Use type hints (where applicable) +5. {"Include comprehensive unit tests" if include_tests else "Focus only on implementation"} +6. {"Include docstrings and comments" if include_docs else "Minimal documentation"} +7. Consider the repository context and maintain consistency with existing code +8. Write code that integrates well with the current branch and recent changes + +**Output Format**: +Please structure your response as follows: + +## Explanation +Brief explanation of your approach and key decisions, considering the repository context. + +## Implementation + +### Main Code +```{language} +# Your main implementation here +``` + +{"### Tests" if include_tests else ""} +{f"```{language}" if include_tests else ""} +{"# Your test code here" if include_tests else ""} +{f"```" if include_tests else ""} + +{"### Documentation" if include_docs else ""} +{"```markdown" if include_docs else ""} +{"# Your documentation here" if include_docs else ""} +{f"```" if include_docs else ""} + +## Commit Message +Suggest a concise git commit message for these changes that follows the repository's commit history style. + +Please ensure the code is production-ready and follows industry standards. +""" + return prompt + + def _parse_response( + self, response: str, language: str, include_tests: bool, include_docs: bool + ) -> Dict[str, Any]: + """Parse Claude's response and extract code files""" + + files = [] + explanation = "" + commit_message = "" + + # Extract explanation + if "## Explanation" in response: + explanation_start = response.find("## Explanation") + len("## Explanation") + explanation_end = response.find("## Implementation") + if explanation_end > explanation_start: + explanation = response[explanation_start:explanation_end].strip() + + # Extract commit message + if "## Commit Message" in response: + commit_start = response.find("## Commit Message") + len("## Commit Message") + commit_message = response[commit_start:].strip() + # Clean up the commit message + commit_message = commit_message.split("\n")[0].strip() + + # Extract main code + main_code = self._extract_code_block(response, "### Main Code", language) + if main_code: + file_ext = self._get_file_extension(language) + files.append( + { + "filename": f"main.{file_ext}", + "content": main_code, + "type": "implementation", + "language": language, + } + ) + + # Extract tests if requested + if include_tests: + test_code = self._extract_code_block(response, "### Tests", language) + if test_code: + test_ext = self._get_file_extension(language) + files.append( + { + "filename": f"test_main.{test_ext}", + "content": test_code, + "type": "test", + "language": language, + } + ) + + # Extract documentation if requested + if include_docs: + docs = self._extract_code_block(response, "### Documentation", "markdown") + if docs: + files.append( + { + "filename": "README.md", + "content": docs, + "type": "documentation", + "language": "markdown", + } + ) + + return { + "files": files, + "explanation": explanation, + "commit_message": commit_message, + } + + def _extract_code_block( + self, text: str, section: str, language: str + ) -> Optional[str]: + """Extract code block from a specific section""" + + section_start = text.find(section) + if section_start == -1: + return None + + # Find the start of the code block + code_start = text.find(f"```{language}", section_start) + if code_start == -1: + code_start = text.find("```", section_start) + if code_start == -1: + return None + + # Find the end of the code block + code_content_start = text.find("\n", code_start) + 1 + code_end = text.find("```", code_content_start) + + if code_end == -1: + return None + + return text[code_content_start:code_end].strip() + + def _get_file_extension(self, language: str) -> str: + """Get appropriate file extension for language""" + extensions = { + "python": "py", + "javascript": "js", + "typescript": "ts", + "java": "java", + "cpp": "cpp", + "c": "c", + "rust": "rs", + "go": "go", + "ruby": "rb", + "php": "php", + "swift": "swift", + "kotlin": "kt", + "scala": "scala", + "r": "R", + "sql": "sql", + "html": "html", + "css": "css", + "shell": "sh", + "bash": "sh", + } + return extensions.get(language.lower(), "txt") + + def _save_files(self, files: List[Dict], tmpdir: str) -> List[str]: + """Save generated files to temporary directory""" + saved_files = [] + + for file_info in files: + file_path = os.path.join(tmpdir, file_info["filename"]) + with open(file_path, "w") as f: + f.write(file_info["content"]) + saved_files.append(file_path) + + return saved_files + + def _run_tests(self, tmpdir: str, language: str) -> Optional[Dict[str, Any]]: + """Run tests for the generated code""" + + try: + if language == "python": + # Try to run pytest + result = subprocess.run( # nosec B603 B607 -- shell=False, static argv; runs pytest on generated code inside an isolated tmpdir + ["python", "-m", "pytest", tmpdir, "-v"], + capture_output=True, + text=True, + timeout=60, + cwd=tmpdir, + ) + + return { + "exit_code": result.returncode, + "stdout": result.stdout, + "stderr": result.stderr, + "success": result.returncode == 0, + } + elif language == "javascript": + # Try to run with node + test_files = [f for f in os.listdir(tmpdir) if f.startswith("test_")] + if test_files: + result = subprocess.run( # nosec B603 B607 -- shell=False, static argv; runs a generated test file inside an isolated tmpdir + ["node", test_files[0]], + capture_output=True, + text=True, + timeout=60, + cwd=tmpdir, + ) + + return { + "exit_code": result.returncode, + "stdout": result.stdout, + "stderr": result.stderr, + "success": result.returncode == 0, + } + + return None + + except subprocess.TimeoutExpired: + return { + "exit_code": -1, + "stdout": "", + "stderr": "Test execution timed out", + "success": False, + } + except Exception as e: + return { + "exit_code": -1, + "stdout": "", + "stderr": f"Test execution error: {str(e)}", + "success": False, + } + + def _build_repository_context( + self, branch_status: Dict[str, Any], commit_history: List[Any] + ) -> str: + """Build repository context string for the prompt""" + context_parts = ["**Repository Context**:"] + + if branch_status: + current_branch = branch_status.get("current_branch", "unknown") + feature_branch = branch_status.get("feature_branch") + has_changes = branch_status.get("has_uncommitted_changes", False) + remote_status = branch_status.get("remote_status", "unknown") + + context_parts.append(f"- Current Branch: `{current_branch}`") + if feature_branch: + context_parts.append(f"- Feature Branch: `{feature_branch}`") + context_parts.append( + f"- Uncommitted Changes: {'Yes' if has_changes else 'No'}" + ) + context_parts.append(f"- Remote Status: {remote_status}") + + if commit_history: + context_parts.append("- Recent Commits:") + for i, commit in enumerate(commit_history[:3]): + context_parts.append(f" {i+1}. `{commit.hash[:8]}` - {commit.message}") + if commit.files_changed: + context_parts.append( + f" Files: {', '.join(commit.files_changed[:5])}" + ) + + if self.repository_context.get("files_changed"): + changed_files = list(set(self.repository_context["files_changed"])) + context_parts.append( + f"- Files Modified This Session: {', '.join(changed_files)}" + ) + + return "\n".join(context_parts) + "\n" + + def _save_files_to_workspace(self, files: List[Dict]) -> List[str]: + """Save generated files directly to workspace""" + saved_files = [] + + for file_info in files: + file_path = os.path.join(self.workspace_path, file_info["filename"]) + + # Create directory if needed + os.makedirs(os.path.dirname(file_path), exist_ok=True) + + with open(file_path, "w") as f: + f.write(file_info["content"]) + saved_files.append(file_path) + + return saved_files + + async def commit_and_push_changes( + self, commit_message: str, files: Optional[List[str]] = None + ) -> Dict[str, Any]: + """Commit and push changes using Git manager""" + if not self.git_manager: + return {"success": False, "error": "No Git manager available"} + + try: + # Commit changes + commit_hash = await self.git_manager.commit_changes( + message=commit_message, + files=files, + iteration_number=self.repository_context["iteration_count"], + ) + + if commit_hash: + self.repository_context["last_commit"] = commit_message + return { + "success": True, + "commit_hash": commit_hash, + "message": "Changes committed successfully", + } + else: + return {"success": True, "message": "No changes to commit"} + + except Exception as e: + return {"success": False, "error": f"Failed to commit changes: {str(e)}"} + + def get_repository_context(self) -> Dict[str, Any]: + """Get current repository context""" + return self.repository_context.copy() + + def reset_context(self): + """Reset the repository context""" + self.repository_context = { + "files_changed": [], + "current_branch": None, + "last_commit": None, + "iteration_count": 0, + } + + async def start_conversation_session(self, sandbox_id: str) -> str: + """Start a conversation session for tracking all Claude Code interactions""" + if not self.agent_id or not self.task_id: + raise ValueError("Agent ID and Task ID required for conversation tracking") + + self.conversation_session_id = ( + await self.conversation_manager.start_conversation_session( + agent_id=self.agent_id, + task_id=self.task_id, + sandbox_id=sandbox_id, + metadata={ + "workspace_path": self.workspace_path, + "model": self.model, + "git_enabled": bool(self.git_manager), + }, + ) + ) + return self.conversation_session_id + + async def end_conversation_session(self) -> bool: + """End the current conversation session""" + if not self.conversation_session_id: + return False + + success = await self.conversation_manager.end_conversation_session( + self.conversation_session_id + ) + self.conversation_session_id = None + return success + + async def execute_task_async( + self, + task: str, + language: str = "python", + context: str = "", + include_tests: bool = True, + include_docs: bool = True, + file_path: Optional[str] = None, + iteration_number: Optional[int] = None, + ) -> Dict[str, Any]: + """Async version of task execution with full conversation tracking""" + + try: + # Update iteration count + if iteration_number: + self.repository_context["iteration_count"] = iteration_number + else: + self.repository_context["iteration_count"] += 1 + + current_iteration = self.repository_context["iteration_count"] + + # Get repository context if Git manager is available + repo_context = "" + if self.git_manager: + try: + branch_status = await self.git_manager.get_branch_status() + self.repository_context["current_branch"] = branch_status.get( + "current_branch" + ) + + commit_history = await self.git_manager.get_commit_history(limit=3) + if commit_history: + self.repository_context["last_commit"] = commit_history[ + 0 + ].message + + repo_context = self._build_repository_context( + branch_status, commit_history + ) + except Exception as e: + repo_context = f"Repository context unavailable: {str(e)}" + + # Read existing file context if provided + existing_code = "" + if file_path: + # Use workspace-relative path if available + full_path = ( + os.path.join(self.workspace_path, file_path) + if not os.path.isabs(file_path) + else file_path + ) + if os.path.exists(full_path): + with open(full_path, "r") as f: + existing_code = f.read() + + # Prepare the comprehensive prompt with repository context + prompt = self._build_prompt( + task=task, + language=language, + context=context, + existing_code=existing_code, + include_tests=include_tests, + include_docs=include_docs, + repo_context=repo_context, + ) + + # Store user prompt in conversation history + if self.conversation_session_id and self.task_id: + await self.conversation_manager.store_user_prompt( + session_id=self.conversation_session_id, + task_id=self.task_id, + iteration_number=current_iteration, + prompt=prompt, + model=self.model, + temperature=0.3, + metadata={ + "task_description": ( + task[:200] + "..." if len(task) > 200 else task + ), + "language": language, + "include_tests": include_tests, + "include_docs": include_docs, + "file_path": file_path, + "workspace_path": self.workspace_path, + }, + ) + + # Record start time for response time tracking + start_time = time.time() + + # Call Claude API + response = self.client.messages.create( + model=self.model, + max_tokens=4096, + temperature=0.3, # Lower temperature for more consistent code + messages=[{"role": "user", "content": prompt}], + ) + + # Extract response content and token usage + response_content = response.content[0].text + token_count = ( + getattr(response.usage, "output_tokens", None) + if hasattr(response, "usage") + else None + ) + + # Store Claude response in conversation history + if self.conversation_session_id and self.task_id: + await self.conversation_manager.store_claude_response( + session_id=self.conversation_session_id, + task_id=self.task_id, + iteration_number=current_iteration, + response=response_content, + token_count=token_count, + model=self.model, + start_time=start_time, + metadata={ + "prompt_length": len(prompt), + "response_length": len(response_content), + }, + ) + + # Parse the response and extract code files + result = self._parse_response( + response_content, language, include_tests, include_docs + ) + + # Save files to workspace if available + saved_files = [] + if self.workspace_path and os.path.exists(self.workspace_path): + saved_files = self._save_files_to_workspace(result["files"]) + + # Store code generations in database + if self.task_id: + for file_info in result["files"]: + await self.conversation_manager.store_code_generation( + task_id=self.task_id, + iteration_number=current_iteration, + file_path=file_info["filename"], + file_type=file_info["type"], + language=file_info.get("language", language), + content=file_info["content"], + ) + + # Update repository context with changed files + self.repository_context["files_changed"].extend( + [ + f["filename"] + for f in result["files"] + if f["type"] == "implementation" + ] + ) + + # Run tests if generated and in workspace + test_results = None + if include_tests and any(f["type"] == "test" for f in result["files"]): + if self.workspace_path and os.path.exists(self.workspace_path): + test_results = self._run_tests(self.workspace_path, language) + + # Store test results + if self.conversation_session_id and self.task_id: + await self.conversation_manager.store_message( + session_id=self.conversation_session_id, + message={ + "task_id": self.task_id, + "iteration_number": current_iteration, + "message_type": MessageType.TEST_RESULT, + "content": json.dumps(test_results), + "metadata": { + "test_framework": ( + "pytest" if language == "python" else "jest" + ), + "workspace_path": self.workspace_path, + }, + }, + ) + + return { + "status": "success", + "files": result["files"], + "saved_files": saved_files, + "explanation": result.get("explanation", ""), + "test_results": test_results, + "commit_message": result.get("commit_message", ""), + "execution_summary": f"Generated {len(result['files'])} files for {language} task: {task[:100]}...", + "iteration": current_iteration, + "workspace_path": self.workspace_path, + "repository_context": self.repository_context, + "conversation_tracked": bool(self.conversation_session_id), + "token_count": token_count, + } + + except Exception as e: + # Store error in conversation history + if self.conversation_session_id and self.task_id: + try: + await self.conversation_manager.store_message( + session_id=self.conversation_session_id, + message={ + "task_id": self.task_id, + "iteration_number": self.repository_context[ + "iteration_count" + ], + "message_type": MessageType.ERROR_MESSAGE, + "content": str(e), + "metadata": { + "error_type": type(e).__name__, + "task_description": ( + task[:200] + "..." if len(task) > 200 else task + ), + }, + }, + ) + except Exception as conv_error: + # Don't let conversation storage errors break the main flow + print( + f"Warning: Failed to store error in conversation: {conv_error}" + ) + + return { + "status": "error", + "error": str(e), + "error_type": "execution_error", + "iteration": self.repository_context["iteration_count"], + "conversation_tracked": bool(self.conversation_session_id), + } diff --git a/services/orchestrator/claude_sdk_manager.py b/services/orchestrator/claude_sdk_manager.py index 3be2939..282da78 100644 --- a/services/orchestrator/claude_sdk_manager.py +++ b/services/orchestrator/claude_sdk_manager.py @@ -1,502 +1,502 @@ -""" -Claude SDK Manager for FuzeAgent - -Manages Claude Code SDK processes, handles interactive states, and integrates -with the File Operations Engine to apply code changes safely. -""" - -import asyncio -import json -import logging -import os -import re -import subprocess # nosec B404 -- used with asyncio.create_subprocess_exec (shell=False) and a static arg list -import time -from dataclasses import dataclass -from datetime import datetime -from enum import Enum -from typing import Any, AsyncGenerator, Callable, Dict, List, Optional - -from .conversation_manager import ConversationManager, MessageType -from .file_operations_engine import FileOperationsEngine, OperationBatch - -logger = logging.getLogger(__name__) - - -class ClaudeSDKState(str, Enum): - IDLE = "idle" - INITIALIZING = "initializing" - RUNNING = "running" - WAITING_FOR_INPUT = "waiting_for_input" - WAITING_FOR_APPROVAL = "waiting_for_approval" - PROCESSING = "processing" - ERROR = "error" - COMPLETED = "completed" - TERMINATED = "terminated" - - -class InteractionType(str, Enum): - USER_INPUT = "user_input" - FILE_APPROVAL = "file_approval" - CONFIRMATION = "confirmation" - SELECTION = "selection" - - -@dataclass -class ClaudeInteraction: - """Represents an interaction request from Claude SDK""" - - interaction_id: str - interaction_type: InteractionType - prompt: str - options: Optional[List[str]] = None - default_response: Optional[str] = None - timeout_seconds: Optional[int] = None - metadata: Optional[Dict[str, Any]] = None - - -@dataclass -class ClaudeSDKSession: - """Represents a Claude SDK session""" - - session_id: str - task_id: str - agent_id: str - workspace_path: str - process: Optional[asyncio.subprocess.Process] = None - state: ClaudeSDKState = ClaudeSDKState.IDLE - current_interaction: Optional[ClaudeInteraction] = None - output_buffer: str = "" - error_buffer: str = "" - started_at: Optional[datetime] = None - last_activity: Optional[datetime] = None - - -class ClaudeSDKManager: - """ - Manages Claude Code SDK processes and handles all interactions. - - Features: - - Interactive process management - - Real-time output streaming - - Human-in-the-loop handling - - File operations integration - - State management and recovery - """ - - def __init__( - self, - file_operations_engine: FileOperationsEngine, - conversation_manager: ConversationManager, - ): - self.file_ops_engine = file_operations_engine - self.conversation_manager = conversation_manager - self.sessions: Dict[str, ClaudeSDKSession] = {} - self.interaction_callbacks: Dict[str, Callable] = {} - - # Configuration - self.claude_cli_path = "claude" # Assume in PATH - self.interaction_timeout = 300 # 5 minutes - self.process_timeout = 3600 # 1 hour - - # Pattern matching for interactive states - self.interaction_patterns = { - InteractionType.USER_INPUT: [ - r"Please provide.*?:", - r"Enter your.*?:", - r"What would you like.*?:", - r"\?\s*$", - ], - InteractionType.FILE_APPROVAL: [ - r"Apply these changes.*?\?", - r"Proceed with.*?file.*?changes.*?\?", - r"Create.*?files.*?\?", - r"Modify.*?files.*?\?", - ], - InteractionType.CONFIRMATION: [ - r"Are you sure.*?\?", - r"Continue.*?\?", - r"Proceed.*?\?", - r"\(y/n\)", - ], - InteractionType.SELECTION: [ - r"Choose.*?:", - r"Select.*?:", - r"\[1\].*?\[2\]", - r"Options.*?:", - ], - } - - async def start_session( - self, - task_id: str, - agent_id: str, - workspace_path: str, - task_description: str, - additional_context: Optional[str] = None, - ) -> str: - """Start a new Claude SDK session""" - - session_id = f"claude-{task_id}-{int(time.time())}" - - session = ClaudeSDKSession( - session_id=session_id, - task_id=task_id, - agent_id=agent_id, - workspace_path=workspace_path, - started_at=datetime.now(), - last_activity=datetime.now(), - ) - - self.sessions[session_id] = session - - try: - # Start Claude Code process - await self._start_claude_process( - session, task_description, additional_context - ) - - # Start output monitoring - asyncio.create_task(self._monitor_session(session)) - - logger.info(f"Started Claude SDK session {session_id}") - return session_id - - except Exception as e: - logger.error(f"Error starting Claude SDK session: {e}") - session.state = ClaudeSDKState.ERROR - raise - - async def send_input(self, session_id: str, user_input: str) -> bool: - """Send input to a Claude SDK session""" - - session = self.sessions.get(session_id) - if not session or not session.process: - return False - - try: - # Send input to process - session.process.stdin.write((user_input + "\n").encode()) - await session.process.stdin.drain() - - # Update session state - session.state = ClaudeSDKState.PROCESSING - session.current_interaction = None - session.last_activity = datetime.now() - - # Store interaction in conversation manager - await self.conversation_manager.store_message( - session_id=session_id, - message={ - "task_id": session.task_id, - "iteration_number": 1, # Would be dynamic in real implementation - "message_type": MessageType.USER_PROMPT, - "content": user_input, - "metadata": {"interaction_type": "human_response"}, - }, - ) - - logger.info(f"Sent input to Claude SDK session {session_id}") - return True - - except Exception as e: - logger.error(f"Error sending input to session {session_id}: {e}") - return False - - async def approve_file_operations( - self, session_id: str, batch_id: str, approved: bool - ) -> bool: - """Approve or reject file operations from Claude SDK""" - - # Apply file operations - success = await self.file_ops_engine.approve_operations(batch_id, approved) - - if success and approved: - # Send approval to Claude SDK - await self.send_input(session_id, "y") - return True - elif success and not approved: - # Send rejection to Claude SDK - await self.send_input(session_id, "n") - return True - - return False - - async def get_session_status(self, session_id: str) -> Optional[Dict[str, Any]]: - """Get current status of a Claude SDK session""" - - session = self.sessions.get(session_id) - if not session: - return None - - return { - "session_id": session_id, - "task_id": session.task_id, - "agent_id": session.agent_id, - "state": session.state.value, - "current_interaction": ( - { - "id": session.current_interaction.interaction_id, - "type": session.current_interaction.interaction_type.value, - "prompt": session.current_interaction.prompt, - "options": session.current_interaction.options, - } - if session.current_interaction - else None - ), - "started_at": ( - session.started_at.isoformat() if session.started_at else None - ), - "last_activity": ( - session.last_activity.isoformat() if session.last_activity else None - ), - "workspace_path": session.workspace_path, - } - - async def terminate_session(self, session_id: str) -> bool: - """Terminate a Claude SDK session""" - - session = self.sessions.get(session_id) - if not session: - return False - - try: - if session.process: - session.process.terminate() - try: - await asyncio.wait_for(session.process.wait(), timeout=10) - except asyncio.TimeoutError: - session.process.kill() - await session.process.wait() - - session.state = ClaudeSDKState.TERMINATED - logger.info(f"Terminated Claude SDK session {session_id}") - return True - - except Exception as e: - logger.error(f"Error terminating session {session_id}: {e}") - return False - - def register_interaction_callback(self, session_id: str, callback: Callable): - """Register callback for interaction events""" - self.interaction_callbacks[session_id] = callback - - async def stream_output(self, session_id: str) -> AsyncGenerator[str, None]: - """Stream real-time output from Claude SDK session""" - - session = self.sessions.get(session_id) - if not session: - return - - last_position = 0 - - while session.state not in [ - ClaudeSDKState.COMPLETED, - ClaudeSDKState.TERMINATED, - ClaudeSDKState.ERROR, - ]: - # Check for new output - if len(session.output_buffer) > last_position: - new_output = session.output_buffer[last_position:] - last_position = len(session.output_buffer) - yield new_output - - await asyncio.sleep(0.1) # Small delay to prevent excessive CPU usage - - # Private methods - - async def _start_claude_process( - self, - session: ClaudeSDKSession, - task_description: str, - additional_context: Optional[str] = None, - ): - """Start the Claude Code CLI process""" - - # Build Claude command - cmd = [ - self.claude_cli_path, - "code", - "--workspace", - session.workspace_path, - "--task", - task_description, - ] - - if additional_context: - cmd.extend(["--context", additional_context]) - - # Set environment - env = os.environ.copy() - env["ANTHROPIC_API_KEY"] = os.environ.get("ANTHROPIC_API_KEY", "") - - # Start process - session.process = await asyncio.create_subprocess_exec( - *cmd, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=session.workspace_path, - env=env, - ) - - session.state = ClaudeSDKState.RUNNING - logger.info(f"Started Claude CLI process for session {session.session_id}") - - async def _monitor_session(self, session: ClaudeSDKSession): - """Monitor a Claude SDK session for output and interactions""" - - logger.info(f"Monitoring Claude SDK session {session.session_id}") - - try: - while session.process and session.process.returncode is None: - # Read output with timeout - try: - output_data = await asyncio.wait_for( - session.process.stdout.read(1024), timeout=0.1 - ) - - if output_data: - output_text = output_data.decode("utf-8", errors="replace") - session.output_buffer += output_text - session.last_activity = datetime.now() - - # Process output for interactions - await self._process_output(session, output_text) - - except asyncio.TimeoutError: - # Check for session timeout - if self._is_session_timed_out(session): - logger.warning(f"Session {session.session_id} timed out") - session.state = ClaudeSDKState.ERROR - await self.terminate_session(session.session_id) - break - - # Small delay to prevent excessive CPU usage - await asyncio.sleep(0.01) - - # Process completed - if session.process and session.process.returncode == 0: - session.state = ClaudeSDKState.COMPLETED - logger.info( - f"Claude SDK session {session.session_id} completed successfully" - ) - else: - session.state = ClaudeSDKState.ERROR - logger.error(f"Claude SDK session {session.session_id} failed") - - except Exception as e: - logger.error(f"Error monitoring session {session.session_id}: {e}") - session.state = ClaudeSDKState.ERROR - - async def _process_output(self, session: ClaudeSDKSession, output_text: str): - """Process output from Claude SDK to detect interactions""" - - # Store output in conversation manager - await self.conversation_manager.store_message( - session_id=session.session_id, - message={ - "task_id": session.task_id, - "iteration_number": 1, # Would be dynamic - "message_type": MessageType.CLAUDE_RESPONSE, - "content": output_text, - "metadata": {"stream_chunk": True}, - }, - ) - - # Check for interaction patterns - interaction = self._detect_interaction(output_text) - if interaction: - session.current_interaction = interaction - session.state = ClaudeSDKState.WAITING_FOR_INPUT - - # Notify callback if registered - callback = self.interaction_callbacks.get(session.session_id) - if callback: - asyncio.create_task(callback(session, interaction)) - - logger.info( - f"Detected interaction in session {session.session_id}: {interaction.interaction_type}" - ) - - # Check for file operations - await self._check_for_file_operations(session, output_text) - - def _detect_interaction(self, output_text: str) -> Optional[ClaudeInteraction]: - """Detect if output contains an interaction request""" - - # Check each interaction type - for interaction_type, patterns in self.interaction_patterns.items(): - for pattern in patterns: - if re.search(pattern, output_text, re.IGNORECASE | re.MULTILINE): - # Extract the prompt (last few lines) - lines = output_text.strip().split("\n") - prompt = "\n".join(lines[-3:]) # Last 3 lines as prompt - - interaction_id = f"interaction-{int(time.time())}" - - return ClaudeInteraction( - interaction_id=interaction_id, - interaction_type=interaction_type, - prompt=prompt, - timeout_seconds=self.interaction_timeout, - ) - - return None - - async def _check_for_file_operations( - self, session: ClaudeSDKSession, output_text: str - ): - """Check if output contains file operation requests""" - - # Look for structured file operations (JSON format) - try: - # Try to find JSON blocks in output - json_blocks = re.findall(r"```json\n(.*?)\n```", output_text, re.DOTALL) - for json_block in json_blocks: - try: - operations_data = json.loads(json_block) - if "operations" in operations_data: - # Process file operations - batch = await self.file_ops_engine.process_claude_response( - operations_data, session.task_id, session.agent_id - ) - - if batch.requires_approval: - session.state = ClaudeSDKState.WAITING_FOR_APPROVAL - session.current_interaction = ClaudeInteraction( - interaction_id=f"approval-{batch.batch_id}", - interaction_type=InteractionType.FILE_APPROVAL, - prompt=f"Approve file operations: {batch.description}", - metadata={"batch_id": batch.batch_id}, - ) - else: - # Auto-approved operations - await self.file_ops_engine.apply_operations_if_approved( - batch.batch_id - ) - - logger.info( - f"Detected file operations in session {session.session_id}" - ) - - except json.JSONDecodeError: - continue - - except Exception as e: - logger.error(f"Error processing file operations: {e}") - - def _is_session_timed_out(self, session: ClaudeSDKSession) -> bool: - """Check if session has timed out""" - - if not session.last_activity: - return False - - timeout_seconds = self.process_timeout - if session.current_interaction: - timeout_seconds = ( - session.current_interaction.timeout_seconds or self.interaction_timeout - ) - - time_since_activity = (datetime.now() - session.last_activity).total_seconds() - return time_since_activity > timeout_seconds +""" +Claude SDK Manager for FuzeAgent + +Manages Claude Code SDK processes, handles interactive states, and integrates +with the File Operations Engine to apply code changes safely. +""" + +import asyncio +import json +import logging +import os +import re +import subprocess # nosec B404 -- used with asyncio.create_subprocess_exec (shell=False) and a static arg list +import time +from dataclasses import dataclass +from datetime import datetime +from enum import Enum +from typing import Any, AsyncGenerator, Callable, Dict, List, Optional + +from .conversation_manager import ConversationManager, MessageType +from .file_operations_engine import FileOperationsEngine, OperationBatch + +logger = logging.getLogger(__name__) + + +class ClaudeSDKState(str, Enum): + IDLE = "idle" + INITIALIZING = "initializing" + RUNNING = "running" + WAITING_FOR_INPUT = "waiting_for_input" + WAITING_FOR_APPROVAL = "waiting_for_approval" + PROCESSING = "processing" + ERROR = "error" + COMPLETED = "completed" + TERMINATED = "terminated" + + +class InteractionType(str, Enum): + USER_INPUT = "user_input" + FILE_APPROVAL = "file_approval" + CONFIRMATION = "confirmation" + SELECTION = "selection" + + +@dataclass +class ClaudeInteraction: + """Represents an interaction request from Claude SDK""" + + interaction_id: str + interaction_type: InteractionType + prompt: str + options: Optional[List[str]] = None + default_response: Optional[str] = None + timeout_seconds: Optional[int] = None + metadata: Optional[Dict[str, Any]] = None + + +@dataclass +class ClaudeSDKSession: + """Represents a Claude SDK session""" + + session_id: str + task_id: str + agent_id: str + workspace_path: str + process: Optional[asyncio.subprocess.Process] = None + state: ClaudeSDKState = ClaudeSDKState.IDLE + current_interaction: Optional[ClaudeInteraction] = None + output_buffer: str = "" + error_buffer: str = "" + started_at: Optional[datetime] = None + last_activity: Optional[datetime] = None + + +class ClaudeSDKManager: + """ + Manages Claude Code SDK processes and handles all interactions. + + Features: + - Interactive process management + - Real-time output streaming + - Human-in-the-loop handling + - File operations integration + - State management and recovery + """ + + def __init__( + self, + file_operations_engine: FileOperationsEngine, + conversation_manager: ConversationManager, + ): + self.file_ops_engine = file_operations_engine + self.conversation_manager = conversation_manager + self.sessions: Dict[str, ClaudeSDKSession] = {} + self.interaction_callbacks: Dict[str, Callable] = {} + + # Configuration + self.claude_cli_path = "claude" # Assume in PATH + self.interaction_timeout = 300 # 5 minutes + self.process_timeout = 3600 # 1 hour + + # Pattern matching for interactive states + self.interaction_patterns = { + InteractionType.USER_INPUT: [ + r"Please provide.*?:", + r"Enter your.*?:", + r"What would you like.*?:", + r"\?\s*$", + ], + InteractionType.FILE_APPROVAL: [ + r"Apply these changes.*?\?", + r"Proceed with.*?file.*?changes.*?\?", + r"Create.*?files.*?\?", + r"Modify.*?files.*?\?", + ], + InteractionType.CONFIRMATION: [ + r"Are you sure.*?\?", + r"Continue.*?\?", + r"Proceed.*?\?", + r"\(y/n\)", + ], + InteractionType.SELECTION: [ + r"Choose.*?:", + r"Select.*?:", + r"\[1\].*?\[2\]", + r"Options.*?:", + ], + } + + async def start_session( + self, + task_id: str, + agent_id: str, + workspace_path: str, + task_description: str, + additional_context: Optional[str] = None, + ) -> str: + """Start a new Claude SDK session""" + + session_id = f"claude-{task_id}-{int(time.time())}" + + session = ClaudeSDKSession( + session_id=session_id, + task_id=task_id, + agent_id=agent_id, + workspace_path=workspace_path, + started_at=datetime.now(), + last_activity=datetime.now(), + ) + + self.sessions[session_id] = session + + try: + # Start Claude Code process + await self._start_claude_process( + session, task_description, additional_context + ) + + # Start output monitoring + asyncio.create_task(self._monitor_session(session)) + + logger.info(f"Started Claude SDK session {session_id}") + return session_id + + except Exception as e: + logger.error(f"Error starting Claude SDK session: {e}") + session.state = ClaudeSDKState.ERROR + raise + + async def send_input(self, session_id: str, user_input: str) -> bool: + """Send input to a Claude SDK session""" + + session = self.sessions.get(session_id) + if not session or not session.process: + return False + + try: + # Send input to process + session.process.stdin.write((user_input + "\n").encode()) + await session.process.stdin.drain() + + # Update session state + session.state = ClaudeSDKState.PROCESSING + session.current_interaction = None + session.last_activity = datetime.now() + + # Store interaction in conversation manager + await self.conversation_manager.store_message( + session_id=session_id, + message={ + "task_id": session.task_id, + "iteration_number": 1, # Would be dynamic in real implementation + "message_type": MessageType.USER_PROMPT, + "content": user_input, + "metadata": {"interaction_type": "human_response"}, + }, + ) + + logger.info(f"Sent input to Claude SDK session {session_id}") + return True + + except Exception as e: + logger.error(f"Error sending input to session {session_id}: {e}") + return False + + async def approve_file_operations( + self, session_id: str, batch_id: str, approved: bool + ) -> bool: + """Approve or reject file operations from Claude SDK""" + + # Apply file operations + success = await self.file_ops_engine.approve_operations(batch_id, approved) + + if success and approved: + # Send approval to Claude SDK + await self.send_input(session_id, "y") + return True + elif success and not approved: + # Send rejection to Claude SDK + await self.send_input(session_id, "n") + return True + + return False + + async def get_session_status(self, session_id: str) -> Optional[Dict[str, Any]]: + """Get current status of a Claude SDK session""" + + session = self.sessions.get(session_id) + if not session: + return None + + return { + "session_id": session_id, + "task_id": session.task_id, + "agent_id": session.agent_id, + "state": session.state.value, + "current_interaction": ( + { + "id": session.current_interaction.interaction_id, + "type": session.current_interaction.interaction_type.value, + "prompt": session.current_interaction.prompt, + "options": session.current_interaction.options, + } + if session.current_interaction + else None + ), + "started_at": ( + session.started_at.isoformat() if session.started_at else None + ), + "last_activity": ( + session.last_activity.isoformat() if session.last_activity else None + ), + "workspace_path": session.workspace_path, + } + + async def terminate_session(self, session_id: str) -> bool: + """Terminate a Claude SDK session""" + + session = self.sessions.get(session_id) + if not session: + return False + + try: + if session.process: + session.process.terminate() + try: + await asyncio.wait_for(session.process.wait(), timeout=10) + except asyncio.TimeoutError: + session.process.kill() + await session.process.wait() + + session.state = ClaudeSDKState.TERMINATED + logger.info(f"Terminated Claude SDK session {session_id}") + return True + + except Exception as e: + logger.error(f"Error terminating session {session_id}: {e}") + return False + + def register_interaction_callback(self, session_id: str, callback: Callable): + """Register callback for interaction events""" + self.interaction_callbacks[session_id] = callback + + async def stream_output(self, session_id: str) -> AsyncGenerator[str, None]: + """Stream real-time output from Claude SDK session""" + + session = self.sessions.get(session_id) + if not session: + return + + last_position = 0 + + while session.state not in [ + ClaudeSDKState.COMPLETED, + ClaudeSDKState.TERMINATED, + ClaudeSDKState.ERROR, + ]: + # Check for new output + if len(session.output_buffer) > last_position: + new_output = session.output_buffer[last_position:] + last_position = len(session.output_buffer) + yield new_output + + await asyncio.sleep(0.1) # Small delay to prevent excessive CPU usage + + # Private methods + + async def _start_claude_process( + self, + session: ClaudeSDKSession, + task_description: str, + additional_context: Optional[str] = None, + ): + """Start the Claude Code CLI process""" + + # Build Claude command + cmd = [ + self.claude_cli_path, + "code", + "--workspace", + session.workspace_path, + "--task", + task_description, + ] + + if additional_context: + cmd.extend(["--context", additional_context]) + + # Set environment + env = os.environ.copy() + env["ANTHROPIC_API_KEY"] = os.environ.get("ANTHROPIC_API_KEY", "") + + # Start process + session.process = await asyncio.create_subprocess_exec( + *cmd, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=session.workspace_path, + env=env, + ) + + session.state = ClaudeSDKState.RUNNING + logger.info(f"Started Claude CLI process for session {session.session_id}") + + async def _monitor_session(self, session: ClaudeSDKSession): + """Monitor a Claude SDK session for output and interactions""" + + logger.info(f"Monitoring Claude SDK session {session.session_id}") + + try: + while session.process and session.process.returncode is None: + # Read output with timeout + try: + output_data = await asyncio.wait_for( + session.process.stdout.read(1024), timeout=0.1 + ) + + if output_data: + output_text = output_data.decode("utf-8", errors="replace") + session.output_buffer += output_text + session.last_activity = datetime.now() + + # Process output for interactions + await self._process_output(session, output_text) + + except asyncio.TimeoutError: + # Check for session timeout + if self._is_session_timed_out(session): + logger.warning(f"Session {session.session_id} timed out") + session.state = ClaudeSDKState.ERROR + await self.terminate_session(session.session_id) + break + + # Small delay to prevent excessive CPU usage + await asyncio.sleep(0.01) + + # Process completed + if session.process and session.process.returncode == 0: + session.state = ClaudeSDKState.COMPLETED + logger.info( + f"Claude SDK session {session.session_id} completed successfully" + ) + else: + session.state = ClaudeSDKState.ERROR + logger.error(f"Claude SDK session {session.session_id} failed") + + except Exception as e: + logger.error(f"Error monitoring session {session.session_id}: {e}") + session.state = ClaudeSDKState.ERROR + + async def _process_output(self, session: ClaudeSDKSession, output_text: str): + """Process output from Claude SDK to detect interactions""" + + # Store output in conversation manager + await self.conversation_manager.store_message( + session_id=session.session_id, + message={ + "task_id": session.task_id, + "iteration_number": 1, # Would be dynamic + "message_type": MessageType.CLAUDE_RESPONSE, + "content": output_text, + "metadata": {"stream_chunk": True}, + }, + ) + + # Check for interaction patterns + interaction = self._detect_interaction(output_text) + if interaction: + session.current_interaction = interaction + session.state = ClaudeSDKState.WAITING_FOR_INPUT + + # Notify callback if registered + callback = self.interaction_callbacks.get(session.session_id) + if callback: + asyncio.create_task(callback(session, interaction)) + + logger.info( + f"Detected interaction in session {session.session_id}: {interaction.interaction_type}" + ) + + # Check for file operations + await self._check_for_file_operations(session, output_text) + + def _detect_interaction(self, output_text: str) -> Optional[ClaudeInteraction]: + """Detect if output contains an interaction request""" + + # Check each interaction type + for interaction_type, patterns in self.interaction_patterns.items(): + for pattern in patterns: + if re.search(pattern, output_text, re.IGNORECASE | re.MULTILINE): + # Extract the prompt (last few lines) + lines = output_text.strip().split("\n") + prompt = "\n".join(lines[-3:]) # Last 3 lines as prompt + + interaction_id = f"interaction-{int(time.time())}" + + return ClaudeInteraction( + interaction_id=interaction_id, + interaction_type=interaction_type, + prompt=prompt, + timeout_seconds=self.interaction_timeout, + ) + + return None + + async def _check_for_file_operations( + self, session: ClaudeSDKSession, output_text: str + ): + """Check if output contains file operation requests""" + + # Look for structured file operations (JSON format) + try: + # Try to find JSON blocks in output + json_blocks = re.findall(r"```json\n(.*?)\n```", output_text, re.DOTALL) + for json_block in json_blocks: + try: + operations_data = json.loads(json_block) + if "operations" in operations_data: + # Process file operations + batch = await self.file_ops_engine.process_claude_response( + operations_data, session.task_id, session.agent_id + ) + + if batch.requires_approval: + session.state = ClaudeSDKState.WAITING_FOR_APPROVAL + session.current_interaction = ClaudeInteraction( + interaction_id=f"approval-{batch.batch_id}", + interaction_type=InteractionType.FILE_APPROVAL, + prompt=f"Approve file operations: {batch.description}", + metadata={"batch_id": batch.batch_id}, + ) + else: + # Auto-approved operations + await self.file_ops_engine.apply_operations_if_approved( + batch.batch_id + ) + + logger.info( + f"Detected file operations in session {session.session_id}" + ) + + except json.JSONDecodeError: + continue + + except Exception as e: + logger.error(f"Error processing file operations: {e}") + + def _is_session_timed_out(self, session: ClaudeSDKSession) -> bool: + """Check if session has timed out""" + + if not session.last_activity: + return False + + timeout_seconds = self.process_timeout + if session.current_interaction: + timeout_seconds = ( + session.current_interaction.timeout_seconds or self.interaction_timeout + ) + + time_since_activity = (datetime.now() - session.last_activity).total_seconds() + return time_since_activity > timeout_seconds diff --git a/services/orchestrator/context_enhancement_service.py b/services/orchestrator/context_enhancement_service.py index bc9ae8b..60e2c9c 100644 --- a/services/orchestrator/context_enhancement_service.py +++ b/services/orchestrator/context_enhancement_service.py @@ -1,774 +1,774 @@ -""" -Context Enhancement Service for FuzeAgent - -This service enhances agent context with relevant organizational and team knowledge -before task execution. It provides intelligent knowledge injection based on -task type, agent capabilities, and historical success patterns. -""" - -import asyncio -import json -import logging -from dataclasses import dataclass -from datetime import datetime -from typing import Any, Dict, List, Optional, Tuple - -import asyncpg - -from .organization_rag_manager import ( - ContentType, - KnowledgeCategory, - KnowledgeSearchResult, - OrganizationRAGManager, -) -from .team_knowledge_manager import TeamKnowledgeManager, TeamKnowledgeSearchResult - -logger = logging.getLogger(__name__) - - -@dataclass -class ContextEnhancement: - """Represents an enhancement to agent context""" - - knowledge_id: str - title: str - content: str - source_type: str # 'organization', 'team', 'agent' - category: str - relevance_score: float - confidence_score: float - usage_stats: Dict[str, Any] - metadata: Dict[str, Any] - - -@dataclass -class EnhancedContext: - """Enhanced context for agent task execution""" - - task_id: str - agent_id: str - team_id: str - organization_id: str - base_context: Dict[str, Any] - organizational_knowledge: List[ContextEnhancement] - team_knowledge: List[ContextEnhancement] - similar_task_insights: List[ContextEnhancement] - success_patterns: List[str] - common_pitfalls: List[str] - recommended_approaches: List[str] - context_summary: str - enhancement_metadata: Dict[str, Any] - - -class ContextEnhancementService: - """ - Enhances agent context with relevant organizational knowledge - to improve task execution success rates. - """ - - def __init__( - self, - database_url: str, - org_rag_manager: OrganizationRAGManager, - team_knowledge_manager: TeamKnowledgeManager, - ): - self.database_url = database_url - self.org_rag_manager = org_rag_manager - self.team_knowledge_manager = team_knowledge_manager - self.pool: Optional[asyncpg.Pool] = None - - # Configuration - self.max_org_knowledge_items = 5 - self.max_team_knowledge_items = 8 - self.max_similar_tasks = 3 - self.min_relevance_threshold = 0.4 - self.context_freshness_days = 90 - - # Statistics - self.enhancements_created = 0 - self.average_enhancement_score = 0.0 - self.knowledge_usage_tracking = {} - - async def initialize(self): - """Initialize the context enhancement service""" - logger.info("Initializing ContextEnhancementService") - - try: - self.pool = await asyncpg.create_pool( - self.database_url, min_size=1, max_size=5, command_timeout=60 - ) - - logger.info("ContextEnhancementService initialized successfully") - - except Exception as e: - logger.error(f"Failed to initialize ContextEnhancementService: {e}") - raise - - async def close(self): - """Close database connections""" - if self.pool: - await self.pool.close() - logger.info("ContextEnhancementService closed") - - async def enhance_agent_context( - self, - agent_id: str, - task_data: Dict[str, Any], - base_context: Optional[Dict[str, Any]] = None, - ) -> EnhancedContext: - """Enhance agent context with relevant organizational knowledge""" - - try: - # Get agent and team information - agent_info = await self._get_agent_info(agent_id) - if not agent_info: - raise ValueError(f"Agent {agent_id} not found") - - # Build search queries based on task data - search_queries = self._build_search_queries(task_data, agent_info) - - # Gather knowledge from different sources - org_knowledge = await self._gather_organizational_knowledge( - agent_info["organization_id"], - search_queries, - agent_id, - agent_info["team_id"], - ) - - team_knowledge = await self._gather_team_knowledge( - agent_info["team_id"], search_queries - ) - - similar_tasks = await self._find_similar_task_insights( - agent_info["organization_id"], task_data, agent_id - ) - - # Extract patterns and recommendations - success_patterns = await self._extract_success_patterns( - org_knowledge + team_knowledge + similar_tasks - ) - - pitfalls = await self._extract_common_pitfalls( - agent_info["organization_id"], task_data - ) - - recommendations = await self._generate_recommendations( - task_data, org_knowledge, team_knowledge, similar_tasks - ) - - # Create context summary - context_summary = self._create_context_summary( - task_data, org_knowledge, team_knowledge, success_patterns - ) - - # Build enhanced context - enhanced_context = EnhancedContext( - task_id=task_data.get("task_id", ""), - agent_id=agent_id, - team_id=agent_info["team_id"], - organization_id=agent_info["organization_id"], - base_context=base_context or {}, - organizational_knowledge=org_knowledge, - team_knowledge=team_knowledge, - similar_task_insights=similar_tasks, - success_patterns=success_patterns, - common_pitfalls=pitfalls, - recommended_approaches=recommendations, - context_summary=context_summary, - enhancement_metadata={ - "enhancement_timestamp": datetime.now().isoformat(), - "search_queries_used": search_queries, - "knowledge_sources_count": { - "organizational": len(org_knowledge), - "team": len(team_knowledge), - "similar_tasks": len(similar_tasks), - }, - "total_relevance_score": sum( - item.relevance_score - for item in org_knowledge + team_knowledge + similar_tasks - ), - "enhancement_version": "1.0", - }, - ) - - # Track enhancement usage - await self._track_enhancement_usage(enhanced_context) - - self.enhancements_created += 1 - - logger.info( - f"Enhanced context for agent {agent_id}: " - f"{len(org_knowledge)} org + {len(team_knowledge)} team + " - f"{len(similar_tasks)} similar task insights" - ) - - return enhanced_context - - except Exception as e: - logger.error(f"Error enhancing context for agent {agent_id}: {e}") - # Return minimal enhanced context on error - return EnhancedContext( - task_id=task_data.get("task_id", ""), - agent_id=agent_id, - team_id="", - organization_id="", - base_context=base_context or {}, - organizational_knowledge=[], - team_knowledge=[], - similar_task_insights=[], - success_patterns=[], - common_pitfalls=[], - recommended_approaches=[], - context_summary="Context enhancement failed - using minimal context", - enhancement_metadata={"error": str(e)}, - ) - - async def get_contextual_guidance( - self, - agent_id: str, - current_task_context: Dict[str, Any], - current_iteration: int = 1, - ) -> Dict[str, Any]: - """Get contextual guidance during task execution""" - - try: - agent_info = await self._get_agent_info(agent_id) - if not agent_info: - return {"guidance": [], "suggestions": []} - - # Build guidance based on current context - guidance_items = [] - - # Get iteration-specific guidance - if current_iteration > 3: - guidance_items.extend( - await self._get_iteration_guidance( - agent_info["organization_id"], current_iteration - ) - ) - - # Get context-specific suggestions - suggestions = await self._get_contextual_suggestions( - agent_info["organization_id"], - agent_info["team_id"], - current_task_context, - ) - - return { - "guidance": guidance_items, - "suggestions": suggestions, - "iteration": current_iteration, - "generated_at": datetime.now().isoformat(), - } - - except Exception as e: - logger.error(f"Error getting contextual guidance: {e}") - return {"guidance": [], "suggestions": []} - - async def update_knowledge_effectiveness( - self, - knowledge_id: str, - knowledge_source: str, - task_success: bool, - agent_feedback: Optional[Dict[str, Any]] = None, - ): - """Update knowledge effectiveness based on usage outcomes""" - - try: - if knowledge_source == "organization": - await self.org_rag_manager.update_knowledge_quality( - knowledge_id=knowledge_id, - success_correlation=1.0 if task_success else -0.2, - feedback_metadata=agent_feedback, - ) - elif knowledge_source == "team": - # Update team knowledge effectiveness - async with self.pool.acquire() as conn: - agent_id = ( - agent_feedback.get("agent_id") if agent_feedback else None - ) - if agent_id: - await self.team_knowledge_manager.update_knowledge_effectiveness( - team_knowledge_id=knowledge_id, - agent_id=agent_id, - task_success=task_success, - feedback_score=agent_feedback.get("usefulness_score"), - usage_context=agent_feedback, - ) - - # Track in local usage statistics - if knowledge_id not in self.knowledge_usage_tracking: - self.knowledge_usage_tracking[knowledge_id] = { - "usage_count": 0, - "success_count": 0, - "effectiveness_score": 0.0, - } - - stats = self.knowledge_usage_tracking[knowledge_id] - stats["usage_count"] += 1 - if task_success: - stats["success_count"] += 1 - stats["effectiveness_score"] = stats["success_count"] / stats["usage_count"] - - except Exception as e: - logger.error(f"Error updating knowledge effectiveness: {e}") - - async def get_enhancement_statistics( - self, - organization_id: Optional[str] = None, - team_id: Optional[str] = None, - days_back: int = 30, - ) -> Dict[str, Any]: - """Get context enhancement statistics""" - - try: - async with self.pool.acquire() as conn: - # Basic enhancement statistics - where_conditions = [ - "created_at >= NOW() - INTERVAL '%s days'" % days_back - ] - params = [] - - if organization_id: - where_conditions.append("organization_id = $1") - params.append(organization_id) - - if team_id: - where_conditions.append( - "team_id = $2" if organization_id else "team_id = $1" - ) - params.append(team_id) - - # This is a placeholder - in practice you'd have a table to track enhancements - stats = { - "total_enhancements": self.enhancements_created, - "average_knowledge_items_per_enhancement": { - "organizational": 3.2, - "team": 4.1, - "similar_tasks": 1.8, - }, - "knowledge_effectiveness": dict(self.knowledge_usage_tracking), - "generated_at": datetime.now().isoformat(), - } - - return stats - - except Exception as e: - logger.error(f"Error getting enhancement statistics: {e}") - return {} - - async def _get_agent_info(self, agent_id: str) -> Optional[Dict[str, Any]]: - """Get agent information including team and organization""" - - async with self.pool.acquire() as conn: - agent_info = await conn.fetchrow( - """ - SELECT a.id, a.name, a.type, a.config, a.team_id, - t.organization_id, t.name as team_name, - o.name as organization_name - FROM agents a - JOIN teams t ON a.team_id = t.id - JOIN organizations o ON t.organization_id = o.id - WHERE a.id = $1 - """, - agent_id, - ) - - return dict(agent_info) if agent_info else None - - def _build_search_queries( - self, task_data: Dict[str, Any], agent_info: Dict[str, Any] - ) -> List[str]: - """Build search queries based on task data and agent information""" - - queries = [] - - # Primary query from task description - if task_data.get("description"): - queries.append(task_data["description"][:200]) - - # Query from task title - if task_data.get("title"): - queries.append(task_data["title"]) - - # Technology-specific queries - if task_data.get("technologies"): - for tech in task_data["technologies"]: - queries.append(f"{tech} development best practices") - - # Task type specific query - if task_data.get("task_type"): - queries.append(f"{task_data['task_type']} implementation guide") - - # Agent type specific query - agent_type = agent_info.get("type", "") - if agent_type: - queries.append(f"{agent_type} workflow best practices") - - return queries[:5] # Limit to top 5 queries - - async def _gather_organizational_knowledge( - self, - organization_id: str, - search_queries: List[str], - agent_id: str, - team_id: str, - ) -> List[ContextEnhancement]: - """Gather relevant organizational knowledge""" - - org_knowledge = [] - - for query in search_queries: - search_results = await self.org_rag_manager.search_knowledge( - organization_id=organization_id, - query=query, - limit=self.max_org_knowledge_items // len(search_queries) + 1, - min_similarity=self.min_relevance_threshold, - requester_agent_id=agent_id, - requester_team_id=team_id, - ) - - for result in search_results: - if result.combined_score >= self.min_relevance_threshold: - enhancement = ContextEnhancement( - knowledge_id=result.knowledge.id, - title=result.knowledge.title, - content=( - result.knowledge.content[:1000] + "..." - if len(result.knowledge.content) > 1000 - else result.knowledge.content - ), - source_type="organization", - category=result.knowledge.knowledge_category.value, - relevance_score=result.combined_score, - confidence_score=result.knowledge.quality_score, - usage_stats={ - "usage_count": result.knowledge.usage_count, - "success_correlation": result.knowledge.success_correlation, - }, - metadata=result.knowledge.metadata, - ) - org_knowledge.append(enhancement) - - # Sort by relevance and remove duplicates - seen_ids = set() - unique_knowledge = [] - for item in sorted( - org_knowledge, key=lambda x: x.relevance_score, reverse=True - ): - if item.knowledge_id not in seen_ids: - unique_knowledge.append(item) - seen_ids.add(item.knowledge_id) - - return unique_knowledge[: self.max_org_knowledge_items] - - async def _gather_team_knowledge( - self, team_id: str, search_queries: List[str] - ) -> List[ContextEnhancement]: - """Gather relevant team knowledge""" - - team_knowledge = [] - - for query in search_queries: - search_results = await self.team_knowledge_manager.search_team_knowledge( - team_id=team_id, - query=query, - limit=self.max_team_knowledge_items // len(search_queries) + 1, - min_similarity=self.min_relevance_threshold, - ) - - for result in search_results: - if result.combined_score >= self.min_relevance_threshold: - enhancement = ContextEnhancement( - knowledge_id=result.team_knowledge.id, - title=result.team_knowledge.title, - content=( - result.team_knowledge.content[:1000] + "..." - if len(result.team_knowledge.content) > 1000 - else result.team_knowledge.content - ), - source_type="team", - category=result.team_knowledge.knowledge_category.value, - relevance_score=result.combined_score, - confidence_score=result.team_knowledge.effectiveness_score, - usage_stats={ - "adoption_rate": result.team_knowledge.agent_adoption_rate, - "effectiveness": result.team_knowledge.effectiveness_score, - }, - metadata=result.team_knowledge.metadata, - ) - team_knowledge.append(enhancement) - - # Sort and deduplicate - seen_ids = set() - unique_knowledge = [] - for item in sorted( - team_knowledge, key=lambda x: x.relevance_score, reverse=True - ): - if item.knowledge_id not in seen_ids: - unique_knowledge.append(item) - seen_ids.add(item.knowledge_id) - - return unique_knowledge[: self.max_team_knowledge_items] - - async def _find_similar_task_insights( - self, organization_id: str, task_data: Dict[str, Any], agent_id: str - ) -> List[ContextEnhancement]: - """Find insights from similar completed tasks""" - - similar_tasks = [] - - try: - async with self.pool.acquire() as conn: - # Find similar tasks based on description similarity and success - similar_task_data = await conn.fetch( - """ - SELECT t.id, t.title, t.description, t.result, t.completed_at, - a.name as agent_name, a.type as agent_type, - similarity(t.description, $2) as similarity_score - FROM tasks t - JOIN agents a ON t.agent_id = a.id - JOIN teams te ON a.team_id = te.id - WHERE te.organization_id = $1 - AND t.status = 'completed' - AND t.result->>'status' = 'completed' - AND t.completed_at >= NOW() - INTERVAL '90 days' - AND similarity(t.description, $2) > 0.3 - ORDER BY similarity_score DESC, t.completed_at DESC - LIMIT $3 - """, - organization_id, - task_data.get("description", ""), - self.max_similar_tasks, - ) - - for task in similar_task_data: - # Extract insights from the task result - task_result = ( - task["result"] if isinstance(task["result"], dict) else {} - ) - - insights_content = self._extract_task_insights( - dict(task), task_result - ) - - if insights_content: - enhancement = ContextEnhancement( - knowledge_id=str(task["id"]), - title=f"Similar Task: {task['title'][:50]}...", - content=insights_content, - source_type="similar_task", - category="process", - relevance_score=float(task["similarity_score"]), - confidence_score=0.8, # High confidence for successful completed tasks - usage_stats={ - "agent_type": task["agent_type"], - "completion_date": task["completed_at"].isoformat(), - }, - metadata={ - "source_task_id": str(task["id"]), - "source_agent": task["agent_name"], - "similarity_score": float(task["similarity_score"]), - }, - ) - similar_tasks.append(enhancement) - - except Exception as e: - logger.error(f"Error finding similar task insights: {e}") - - return similar_tasks - - async def _extract_success_patterns( - self, all_knowledge: List[ContextEnhancement] - ) -> List[str]: - """Extract success patterns from knowledge items""" - - patterns = [] - - for item in all_knowledge: - # Look for success indicators in metadata - if "success_indicators" in item.metadata: - patterns.extend(item.metadata["success_indicators"]) - - # Extract patterns from high-confidence, high-usage items - if item.confidence_score > 0.7 and item.relevance_score > 0.6: - if "optimization" in item.title.lower(): - patterns.append("Focus on optimization early") - if "test" in item.title.lower(): - patterns.append("Comprehensive testing leads to success") - if "pattern" in item.title.lower(): - patterns.append("Follow established patterns") - - return list(set(patterns)) # Remove duplicates - - async def _extract_common_pitfalls( - self, organization_id: str, task_data: Dict[str, Any] - ) -> List[str]: - """Extract common pitfalls for this type of task""" - - pitfalls = [] - - try: - # Search for error patterns and failure knowledge - error_knowledge = await self.org_rag_manager.search_knowledge( - organization_id=organization_id, - query=f"error pattern {task_data.get('task_type', '')}", - categories=[KnowledgeCategory.TROUBLESHOOTING], - limit=5, - min_similarity=0.3, - ) - - for result in error_knowledge: - if "failure_patterns" in result.knowledge.metadata: - pitfalls.extend(result.knowledge.metadata["failure_patterns"]) - - # Extract pitfalls from error pattern content - content_lower = result.knowledge.content.lower() - if ( - "avoid" in content_lower - or "pitfall" in content_lower - or "common mistake" in content_lower - ): - pitfalls.append(result.knowledge.title) - - except Exception as e: - logger.error(f"Error extracting pitfalls: {e}") - - return list(set(pitfalls))[:5] # Top 5 pitfalls - - async def _generate_recommendations( - self, - task_data: Dict[str, Any], - org_knowledge: List[ContextEnhancement], - team_knowledge: List[ContextEnhancement], - similar_tasks: List[ContextEnhancement], - ) -> List[str]: - """Generate actionable recommendations based on knowledge""" - - recommendations = [] - - # Recommendations from high-value organizational knowledge - high_value_org = [item for item in org_knowledge if item.confidence_score > 0.7] - for item in high_value_org[:3]: - if item.category == "best_practice": - recommendations.append(f"Apply best practice: {item.title}") - elif item.category == "development": - recommendations.append(f"Consider development approach: {item.title}") - - # Recommendations from effective team knowledge - effective_team = [ - item - for item in team_knowledge - if item.usage_stats.get("adoption_rate", 0) > 0.5 - ] - for item in effective_team[:2]: - recommendations.append(f"Team recommendation: {item.title}") - - # Recommendations from similar successful tasks - for task in similar_tasks: - if task.relevance_score > 0.6: - recommendations.append( - f"Based on similar task: Consider approach used in '{task.title}'" - ) - - return recommendations[:8] # Limit recommendations - - def _create_context_summary( - self, - task_data: Dict[str, Any], - org_knowledge: List[ContextEnhancement], - team_knowledge: List[ContextEnhancement], - success_patterns: List[str], - ) -> str: - """Create a summary of the enhanced context""" - - summary_parts = [] - - summary_parts.append(f"Enhanced context for: {task_data.get('title', 'Task')}") - - if org_knowledge: - summary_parts.append( - f"• {len(org_knowledge)} organizational knowledge items available" - ) - - if team_knowledge: - summary_parts.append( - f"• {len(team_knowledge)} team-specific insights included" - ) - - if success_patterns: - summary_parts.append( - f"• {len(success_patterns)} success patterns identified" - ) - summary_parts.append(f"Key patterns: {', '.join(success_patterns[:3])}") - - return "\n".join(summary_parts) - - def _extract_task_insights(self, task_data: Dict, task_result: Dict) -> str: - """Extract insights from a completed task""" - - insights = [] - - # Extract approach information - if task_result.get("iterations"): - insights.append(f"Completed in {task_result['iterations']} iterations") - - if task_result.get("pull_request_url"): - insights.append("Successfully created pull request") - - # Extract process information - if task_data.get("description"): - insights.append(f"Approach: {task_data['description'][:100]}...") - - return "\n".join(insights) - - async def _get_iteration_guidance( - self, organization_id: str, iteration_count: int - ) -> List[str]: - """Get guidance for high iteration count situations""" - - guidance = [] - - if iteration_count > 5: - # Search for guidance on complex tasks - complex_task_knowledge = await self.org_rag_manager.search_knowledge( - organization_id=organization_id, - query="complex task multiple iterations debugging", - limit=3, - min_similarity=0.3, - ) - - for result in complex_task_knowledge: - if "process" in result.knowledge.knowledge_category.value: - guidance.append(f"Process guidance: {result.knowledge.title}") - - return guidance - - async def _get_contextual_suggestions( - self, organization_id: str, team_id: str, current_context: Dict[str, Any] - ) -> List[str]: - """Get suggestions based on current execution context""" - - suggestions = [] - - # Context-specific suggestions based on current state - if current_context.get("error_count", 0) > 2: - suggestions.append( - "Consider reviewing error patterns in organizational knowledge" - ) - - if current_context.get("execution_time_minutes", 0) > 60: - suggestions.append("Look for optimization guidance from team knowledge") - - return suggestions - - async def _track_enhancement_usage(self, enhanced_context: EnhancedContext): - """Track usage of enhancement for analytics""" - - try: - async with self.pool.acquire() as conn: - # This would store enhancement usage data for analytics - # Placeholder for actual implementation - pass - except Exception as e: - logger.error(f"Error tracking enhancement usage: {e}") +""" +Context Enhancement Service for FuzeAgent + +This service enhances agent context with relevant organizational and team knowledge +before task execution. It provides intelligent knowledge injection based on +task type, agent capabilities, and historical success patterns. +""" + +import asyncio +import json +import logging +from dataclasses import dataclass +from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple + +import asyncpg + +from .organization_rag_manager import ( + ContentType, + KnowledgeCategory, + KnowledgeSearchResult, + OrganizationRAGManager, +) +from .team_knowledge_manager import TeamKnowledgeManager, TeamKnowledgeSearchResult + +logger = logging.getLogger(__name__) + + +@dataclass +class ContextEnhancement: + """Represents an enhancement to agent context""" + + knowledge_id: str + title: str + content: str + source_type: str # 'organization', 'team', 'agent' + category: str + relevance_score: float + confidence_score: float + usage_stats: Dict[str, Any] + metadata: Dict[str, Any] + + +@dataclass +class EnhancedContext: + """Enhanced context for agent task execution""" + + task_id: str + agent_id: str + team_id: str + organization_id: str + base_context: Dict[str, Any] + organizational_knowledge: List[ContextEnhancement] + team_knowledge: List[ContextEnhancement] + similar_task_insights: List[ContextEnhancement] + success_patterns: List[str] + common_pitfalls: List[str] + recommended_approaches: List[str] + context_summary: str + enhancement_metadata: Dict[str, Any] + + +class ContextEnhancementService: + """ + Enhances agent context with relevant organizational knowledge + to improve task execution success rates. + """ + + def __init__( + self, + database_url: str, + org_rag_manager: OrganizationRAGManager, + team_knowledge_manager: TeamKnowledgeManager, + ): + self.database_url = database_url + self.org_rag_manager = org_rag_manager + self.team_knowledge_manager = team_knowledge_manager + self.pool: Optional[asyncpg.Pool] = None + + # Configuration + self.max_org_knowledge_items = 5 + self.max_team_knowledge_items = 8 + self.max_similar_tasks = 3 + self.min_relevance_threshold = 0.4 + self.context_freshness_days = 90 + + # Statistics + self.enhancements_created = 0 + self.average_enhancement_score = 0.0 + self.knowledge_usage_tracking = {} + + async def initialize(self): + """Initialize the context enhancement service""" + logger.info("Initializing ContextEnhancementService") + + try: + self.pool = await asyncpg.create_pool( + self.database_url, min_size=1, max_size=5, command_timeout=60 + ) + + logger.info("ContextEnhancementService initialized successfully") + + except Exception as e: + logger.error(f"Failed to initialize ContextEnhancementService: {e}") + raise + + async def close(self): + """Close database connections""" + if self.pool: + await self.pool.close() + logger.info("ContextEnhancementService closed") + + async def enhance_agent_context( + self, + agent_id: str, + task_data: Dict[str, Any], + base_context: Optional[Dict[str, Any]] = None, + ) -> EnhancedContext: + """Enhance agent context with relevant organizational knowledge""" + + try: + # Get agent and team information + agent_info = await self._get_agent_info(agent_id) + if not agent_info: + raise ValueError(f"Agent {agent_id} not found") + + # Build search queries based on task data + search_queries = self._build_search_queries(task_data, agent_info) + + # Gather knowledge from different sources + org_knowledge = await self._gather_organizational_knowledge( + agent_info["organization_id"], + search_queries, + agent_id, + agent_info["team_id"], + ) + + team_knowledge = await self._gather_team_knowledge( + agent_info["team_id"], search_queries + ) + + similar_tasks = await self._find_similar_task_insights( + agent_info["organization_id"], task_data, agent_id + ) + + # Extract patterns and recommendations + success_patterns = await self._extract_success_patterns( + org_knowledge + team_knowledge + similar_tasks + ) + + pitfalls = await self._extract_common_pitfalls( + agent_info["organization_id"], task_data + ) + + recommendations = await self._generate_recommendations( + task_data, org_knowledge, team_knowledge, similar_tasks + ) + + # Create context summary + context_summary = self._create_context_summary( + task_data, org_knowledge, team_knowledge, success_patterns + ) + + # Build enhanced context + enhanced_context = EnhancedContext( + task_id=task_data.get("task_id", ""), + agent_id=agent_id, + team_id=agent_info["team_id"], + organization_id=agent_info["organization_id"], + base_context=base_context or {}, + organizational_knowledge=org_knowledge, + team_knowledge=team_knowledge, + similar_task_insights=similar_tasks, + success_patterns=success_patterns, + common_pitfalls=pitfalls, + recommended_approaches=recommendations, + context_summary=context_summary, + enhancement_metadata={ + "enhancement_timestamp": datetime.now().isoformat(), + "search_queries_used": search_queries, + "knowledge_sources_count": { + "organizational": len(org_knowledge), + "team": len(team_knowledge), + "similar_tasks": len(similar_tasks), + }, + "total_relevance_score": sum( + item.relevance_score + for item in org_knowledge + team_knowledge + similar_tasks + ), + "enhancement_version": "1.0", + }, + ) + + # Track enhancement usage + await self._track_enhancement_usage(enhanced_context) + + self.enhancements_created += 1 + + logger.info( + f"Enhanced context for agent {agent_id}: " + f"{len(org_knowledge)} org + {len(team_knowledge)} team + " + f"{len(similar_tasks)} similar task insights" + ) + + return enhanced_context + + except Exception as e: + logger.error(f"Error enhancing context for agent {agent_id}: {e}") + # Return minimal enhanced context on error + return EnhancedContext( + task_id=task_data.get("task_id", ""), + agent_id=agent_id, + team_id="", + organization_id="", + base_context=base_context or {}, + organizational_knowledge=[], + team_knowledge=[], + similar_task_insights=[], + success_patterns=[], + common_pitfalls=[], + recommended_approaches=[], + context_summary="Context enhancement failed - using minimal context", + enhancement_metadata={"error": str(e)}, + ) + + async def get_contextual_guidance( + self, + agent_id: str, + current_task_context: Dict[str, Any], + current_iteration: int = 1, + ) -> Dict[str, Any]: + """Get contextual guidance during task execution""" + + try: + agent_info = await self._get_agent_info(agent_id) + if not agent_info: + return {"guidance": [], "suggestions": []} + + # Build guidance based on current context + guidance_items = [] + + # Get iteration-specific guidance + if current_iteration > 3: + guidance_items.extend( + await self._get_iteration_guidance( + agent_info["organization_id"], current_iteration + ) + ) + + # Get context-specific suggestions + suggestions = await self._get_contextual_suggestions( + agent_info["organization_id"], + agent_info["team_id"], + current_task_context, + ) + + return { + "guidance": guidance_items, + "suggestions": suggestions, + "iteration": current_iteration, + "generated_at": datetime.now().isoformat(), + } + + except Exception as e: + logger.error(f"Error getting contextual guidance: {e}") + return {"guidance": [], "suggestions": []} + + async def update_knowledge_effectiveness( + self, + knowledge_id: str, + knowledge_source: str, + task_success: bool, + agent_feedback: Optional[Dict[str, Any]] = None, + ): + """Update knowledge effectiveness based on usage outcomes""" + + try: + if knowledge_source == "organization": + await self.org_rag_manager.update_knowledge_quality( + knowledge_id=knowledge_id, + success_correlation=1.0 if task_success else -0.2, + feedback_metadata=agent_feedback, + ) + elif knowledge_source == "team": + # Update team knowledge effectiveness + async with self.pool.acquire() as conn: + agent_id = ( + agent_feedback.get("agent_id") if agent_feedback else None + ) + if agent_id: + await self.team_knowledge_manager.update_knowledge_effectiveness( + team_knowledge_id=knowledge_id, + agent_id=agent_id, + task_success=task_success, + feedback_score=agent_feedback.get("usefulness_score"), + usage_context=agent_feedback, + ) + + # Track in local usage statistics + if knowledge_id not in self.knowledge_usage_tracking: + self.knowledge_usage_tracking[knowledge_id] = { + "usage_count": 0, + "success_count": 0, + "effectiveness_score": 0.0, + } + + stats = self.knowledge_usage_tracking[knowledge_id] + stats["usage_count"] += 1 + if task_success: + stats["success_count"] += 1 + stats["effectiveness_score"] = stats["success_count"] / stats["usage_count"] + + except Exception as e: + logger.error(f"Error updating knowledge effectiveness: {e}") + + async def get_enhancement_statistics( + self, + organization_id: Optional[str] = None, + team_id: Optional[str] = None, + days_back: int = 30, + ) -> Dict[str, Any]: + """Get context enhancement statistics""" + + try: + async with self.pool.acquire() as conn: + # Basic enhancement statistics + where_conditions = [ + "created_at >= NOW() - INTERVAL '%s days'" % days_back + ] + params = [] + + if organization_id: + where_conditions.append("organization_id = $1") + params.append(organization_id) + + if team_id: + where_conditions.append( + "team_id = $2" if organization_id else "team_id = $1" + ) + params.append(team_id) + + # This is a placeholder - in practice you'd have a table to track enhancements + stats = { + "total_enhancements": self.enhancements_created, + "average_knowledge_items_per_enhancement": { + "organizational": 3.2, + "team": 4.1, + "similar_tasks": 1.8, + }, + "knowledge_effectiveness": dict(self.knowledge_usage_tracking), + "generated_at": datetime.now().isoformat(), + } + + return stats + + except Exception as e: + logger.error(f"Error getting enhancement statistics: {e}") + return {} + + async def _get_agent_info(self, agent_id: str) -> Optional[Dict[str, Any]]: + """Get agent information including team and organization""" + + async with self.pool.acquire() as conn: + agent_info = await conn.fetchrow( + """ + SELECT a.id, a.name, a.type, a.config, a.team_id, + t.organization_id, t.name as team_name, + o.name as organization_name + FROM agents a + JOIN teams t ON a.team_id = t.id + JOIN organizations o ON t.organization_id = o.id + WHERE a.id = $1 + """, + agent_id, + ) + + return dict(agent_info) if agent_info else None + + def _build_search_queries( + self, task_data: Dict[str, Any], agent_info: Dict[str, Any] + ) -> List[str]: + """Build search queries based on task data and agent information""" + + queries = [] + + # Primary query from task description + if task_data.get("description"): + queries.append(task_data["description"][:200]) + + # Query from task title + if task_data.get("title"): + queries.append(task_data["title"]) + + # Technology-specific queries + if task_data.get("technologies"): + for tech in task_data["technologies"]: + queries.append(f"{tech} development best practices") + + # Task type specific query + if task_data.get("task_type"): + queries.append(f"{task_data['task_type']} implementation guide") + + # Agent type specific query + agent_type = agent_info.get("type", "") + if agent_type: + queries.append(f"{agent_type} workflow best practices") + + return queries[:5] # Limit to top 5 queries + + async def _gather_organizational_knowledge( + self, + organization_id: str, + search_queries: List[str], + agent_id: str, + team_id: str, + ) -> List[ContextEnhancement]: + """Gather relevant organizational knowledge""" + + org_knowledge = [] + + for query in search_queries: + search_results = await self.org_rag_manager.search_knowledge( + organization_id=organization_id, + query=query, + limit=self.max_org_knowledge_items // len(search_queries) + 1, + min_similarity=self.min_relevance_threshold, + requester_agent_id=agent_id, + requester_team_id=team_id, + ) + + for result in search_results: + if result.combined_score >= self.min_relevance_threshold: + enhancement = ContextEnhancement( + knowledge_id=result.knowledge.id, + title=result.knowledge.title, + content=( + result.knowledge.content[:1000] + "..." + if len(result.knowledge.content) > 1000 + else result.knowledge.content + ), + source_type="organization", + category=result.knowledge.knowledge_category.value, + relevance_score=result.combined_score, + confidence_score=result.knowledge.quality_score, + usage_stats={ + "usage_count": result.knowledge.usage_count, + "success_correlation": result.knowledge.success_correlation, + }, + metadata=result.knowledge.metadata, + ) + org_knowledge.append(enhancement) + + # Sort by relevance and remove duplicates + seen_ids = set() + unique_knowledge = [] + for item in sorted( + org_knowledge, key=lambda x: x.relevance_score, reverse=True + ): + if item.knowledge_id not in seen_ids: + unique_knowledge.append(item) + seen_ids.add(item.knowledge_id) + + return unique_knowledge[: self.max_org_knowledge_items] + + async def _gather_team_knowledge( + self, team_id: str, search_queries: List[str] + ) -> List[ContextEnhancement]: + """Gather relevant team knowledge""" + + team_knowledge = [] + + for query in search_queries: + search_results = await self.team_knowledge_manager.search_team_knowledge( + team_id=team_id, + query=query, + limit=self.max_team_knowledge_items // len(search_queries) + 1, + min_similarity=self.min_relevance_threshold, + ) + + for result in search_results: + if result.combined_score >= self.min_relevance_threshold: + enhancement = ContextEnhancement( + knowledge_id=result.team_knowledge.id, + title=result.team_knowledge.title, + content=( + result.team_knowledge.content[:1000] + "..." + if len(result.team_knowledge.content) > 1000 + else result.team_knowledge.content + ), + source_type="team", + category=result.team_knowledge.knowledge_category.value, + relevance_score=result.combined_score, + confidence_score=result.team_knowledge.effectiveness_score, + usage_stats={ + "adoption_rate": result.team_knowledge.agent_adoption_rate, + "effectiveness": result.team_knowledge.effectiveness_score, + }, + metadata=result.team_knowledge.metadata, + ) + team_knowledge.append(enhancement) + + # Sort and deduplicate + seen_ids = set() + unique_knowledge = [] + for item in sorted( + team_knowledge, key=lambda x: x.relevance_score, reverse=True + ): + if item.knowledge_id not in seen_ids: + unique_knowledge.append(item) + seen_ids.add(item.knowledge_id) + + return unique_knowledge[: self.max_team_knowledge_items] + + async def _find_similar_task_insights( + self, organization_id: str, task_data: Dict[str, Any], agent_id: str + ) -> List[ContextEnhancement]: + """Find insights from similar completed tasks""" + + similar_tasks = [] + + try: + async with self.pool.acquire() as conn: + # Find similar tasks based on description similarity and success + similar_task_data = await conn.fetch( + """ + SELECT t.id, t.title, t.description, t.result, t.completed_at, + a.name as agent_name, a.type as agent_type, + similarity(t.description, $2) as similarity_score + FROM tasks t + JOIN agents a ON t.agent_id = a.id + JOIN teams te ON a.team_id = te.id + WHERE te.organization_id = $1 + AND t.status = 'completed' + AND t.result->>'status' = 'completed' + AND t.completed_at >= NOW() - INTERVAL '90 days' + AND similarity(t.description, $2) > 0.3 + ORDER BY similarity_score DESC, t.completed_at DESC + LIMIT $3 + """, + organization_id, + task_data.get("description", ""), + self.max_similar_tasks, + ) + + for task in similar_task_data: + # Extract insights from the task result + task_result = ( + task["result"] if isinstance(task["result"], dict) else {} + ) + + insights_content = self._extract_task_insights( + dict(task), task_result + ) + + if insights_content: + enhancement = ContextEnhancement( + knowledge_id=str(task["id"]), + title=f"Similar Task: {task['title'][:50]}...", + content=insights_content, + source_type="similar_task", + category="process", + relevance_score=float(task["similarity_score"]), + confidence_score=0.8, # High confidence for successful completed tasks + usage_stats={ + "agent_type": task["agent_type"], + "completion_date": task["completed_at"].isoformat(), + }, + metadata={ + "source_task_id": str(task["id"]), + "source_agent": task["agent_name"], + "similarity_score": float(task["similarity_score"]), + }, + ) + similar_tasks.append(enhancement) + + except Exception as e: + logger.error(f"Error finding similar task insights: {e}") + + return similar_tasks + + async def _extract_success_patterns( + self, all_knowledge: List[ContextEnhancement] + ) -> List[str]: + """Extract success patterns from knowledge items""" + + patterns = [] + + for item in all_knowledge: + # Look for success indicators in metadata + if "success_indicators" in item.metadata: + patterns.extend(item.metadata["success_indicators"]) + + # Extract patterns from high-confidence, high-usage items + if item.confidence_score > 0.7 and item.relevance_score > 0.6: + if "optimization" in item.title.lower(): + patterns.append("Focus on optimization early") + if "test" in item.title.lower(): + patterns.append("Comprehensive testing leads to success") + if "pattern" in item.title.lower(): + patterns.append("Follow established patterns") + + return list(set(patterns)) # Remove duplicates + + async def _extract_common_pitfalls( + self, organization_id: str, task_data: Dict[str, Any] + ) -> List[str]: + """Extract common pitfalls for this type of task""" + + pitfalls = [] + + try: + # Search for error patterns and failure knowledge + error_knowledge = await self.org_rag_manager.search_knowledge( + organization_id=organization_id, + query=f"error pattern {task_data.get('task_type', '')}", + categories=[KnowledgeCategory.TROUBLESHOOTING], + limit=5, + min_similarity=0.3, + ) + + for result in error_knowledge: + if "failure_patterns" in result.knowledge.metadata: + pitfalls.extend(result.knowledge.metadata["failure_patterns"]) + + # Extract pitfalls from error pattern content + content_lower = result.knowledge.content.lower() + if ( + "avoid" in content_lower + or "pitfall" in content_lower + or "common mistake" in content_lower + ): + pitfalls.append(result.knowledge.title) + + except Exception as e: + logger.error(f"Error extracting pitfalls: {e}") + + return list(set(pitfalls))[:5] # Top 5 pitfalls + + async def _generate_recommendations( + self, + task_data: Dict[str, Any], + org_knowledge: List[ContextEnhancement], + team_knowledge: List[ContextEnhancement], + similar_tasks: List[ContextEnhancement], + ) -> List[str]: + """Generate actionable recommendations based on knowledge""" + + recommendations = [] + + # Recommendations from high-value organizational knowledge + high_value_org = [item for item in org_knowledge if item.confidence_score > 0.7] + for item in high_value_org[:3]: + if item.category == "best_practice": + recommendations.append(f"Apply best practice: {item.title}") + elif item.category == "development": + recommendations.append(f"Consider development approach: {item.title}") + + # Recommendations from effective team knowledge + effective_team = [ + item + for item in team_knowledge + if item.usage_stats.get("adoption_rate", 0) > 0.5 + ] + for item in effective_team[:2]: + recommendations.append(f"Team recommendation: {item.title}") + + # Recommendations from similar successful tasks + for task in similar_tasks: + if task.relevance_score > 0.6: + recommendations.append( + f"Based on similar task: Consider approach used in '{task.title}'" + ) + + return recommendations[:8] # Limit recommendations + + def _create_context_summary( + self, + task_data: Dict[str, Any], + org_knowledge: List[ContextEnhancement], + team_knowledge: List[ContextEnhancement], + success_patterns: List[str], + ) -> str: + """Create a summary of the enhanced context""" + + summary_parts = [] + + summary_parts.append(f"Enhanced context for: {task_data.get('title', 'Task')}") + + if org_knowledge: + summary_parts.append( + f"• {len(org_knowledge)} organizational knowledge items available" + ) + + if team_knowledge: + summary_parts.append( + f"• {len(team_knowledge)} team-specific insights included" + ) + + if success_patterns: + summary_parts.append( + f"• {len(success_patterns)} success patterns identified" + ) + summary_parts.append(f"Key patterns: {', '.join(success_patterns[:3])}") + + return "\n".join(summary_parts) + + def _extract_task_insights(self, task_data: Dict, task_result: Dict) -> str: + """Extract insights from a completed task""" + + insights = [] + + # Extract approach information + if task_result.get("iterations"): + insights.append(f"Completed in {task_result['iterations']} iterations") + + if task_result.get("pull_request_url"): + insights.append("Successfully created pull request") + + # Extract process information + if task_data.get("description"): + insights.append(f"Approach: {task_data['description'][:100]}...") + + return "\n".join(insights) + + async def _get_iteration_guidance( + self, organization_id: str, iteration_count: int + ) -> List[str]: + """Get guidance for high iteration count situations""" + + guidance = [] + + if iteration_count > 5: + # Search for guidance on complex tasks + complex_task_knowledge = await self.org_rag_manager.search_knowledge( + organization_id=organization_id, + query="complex task multiple iterations debugging", + limit=3, + min_similarity=0.3, + ) + + for result in complex_task_knowledge: + if "process" in result.knowledge.knowledge_category.value: + guidance.append(f"Process guidance: {result.knowledge.title}") + + return guidance + + async def _get_contextual_suggestions( + self, organization_id: str, team_id: str, current_context: Dict[str, Any] + ) -> List[str]: + """Get suggestions based on current execution context""" + + suggestions = [] + + # Context-specific suggestions based on current state + if current_context.get("error_count", 0) > 2: + suggestions.append( + "Consider reviewing error patterns in organizational knowledge" + ) + + if current_context.get("execution_time_minutes", 0) > 60: + suggestions.append("Look for optimization guidance from team knowledge") + + return suggestions + + async def _track_enhancement_usage(self, enhanced_context: EnhancedContext): + """Track usage of enhancement for analytics""" + + try: + async with self.pool.acquire() as conn: + # This would store enhancement usage data for analytics + # Placeholder for actual implementation + pass + except Exception as e: + logger.error(f"Error tracking enhancement usage: {e}") diff --git a/services/orchestrator/context_service.py b/services/orchestrator/context_service.py index e7092b4..7dd7bab 100644 --- a/services/orchestrator/context_service.py +++ b/services/orchestrator/context_service.py @@ -1,152 +1,152 @@ -import asyncio -import json -from typing import Dict, List, Optional - -import numpy as np -from sentence_transformers import SentenceTransformer - -from .database import get_db_connection - - -class ContextService: - def __init__(self): - # Load sentence transformer for embeddings - self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") - - def generate_embedding(self, text: str) -> List[float]: - """Generate embedding for text""" - embedding = self.embedding_model.encode(text) - return embedding.tolist() - - async def store_interaction( - self, agent_id: str, content: str, metadata: Dict = None - ) -> str: - """Store agent interaction with embedding""" - embedding = self.generate_embedding(content) - - async with get_db_connection() as conn: - interaction_id = await conn.fetchval( - """ - INSERT INTO interactions (agent_id, content, embedding, metadata) - VALUES ($1, $2, $3, $4) - RETURNING id - """, - agent_id, - content, - embedding, - metadata or {}, - ) - return str(interaction_id) - - async def get_similar_interactions( - self, - query: str, - agent_id: str = None, - limit: int = 5, - similarity_threshold: float = 0.7, - ) -> List[Dict]: - """Find similar interactions using vector similarity""" - query_embedding = self.generate_embedding(query) - - async with get_db_connection() as conn: - if agent_id: - rows = await conn.fetch( - """ - SELECT id, agent_id, content, metadata, created_at, - 1 - (embedding <=> $1) as similarity - FROM interactions - WHERE agent_id = $2 AND 1 - (embedding <=> $1) > $3 - ORDER BY similarity DESC - LIMIT $4 - """, - query_embedding, - agent_id, - similarity_threshold, - limit, - ) - else: - rows = await conn.fetch( - """ - SELECT id, agent_id, content, metadata, created_at, - 1 - (embedding <=> $1) as similarity - FROM interactions - WHERE 1 - (embedding <=> $1) > $2 - ORDER BY similarity DESC - LIMIT $3 - """, - query_embedding, - similarity_threshold, - limit, - ) - - return [dict(row) for row in rows] - - async def get_context(self, query: str, agent_id: str = None) -> Dict: - """Get relevant context for a query""" - similar_interactions = await self.get_similar_interactions(query, agent_id) - - # Get recent interactions from same agent - recent_interactions = [] - if agent_id: - async with get_db_connection() as conn: - rows = await conn.fetch( - """ - SELECT content, metadata, created_at - FROM interactions - WHERE agent_id = $1 - ORDER BY created_at DESC - LIMIT 10 - """, - agent_id, - ) - recent_interactions = [dict(row) for row in rows] - - return { - "similar_interactions": similar_interactions, - "recent_interactions": recent_interactions, - "relevant_code": self._extract_code_snippets(similar_interactions), - "similar_features": self._extract_similar_features(similar_interactions), - } - - def _extract_code_snippets(self, interactions: List[Dict]) -> str: - """Extract code snippets from interactions""" - code_snippets = [] - for interaction in interactions: - content = interaction.get("content", "") - # Simple extraction - look for code blocks - if "```" in content: - parts = content.split("```") - for i in range(1, len(parts), 2): - code_snippets.append(parts[i].strip()) - - return "\n\n".join(code_snippets[:3]) # Return top 3 snippets - - def _extract_similar_features(self, interactions: List[Dict]) -> str: - """Extract similar feature descriptions""" - features = [] - for interaction in interactions: - metadata = interaction.get("metadata", {}) - if "task_type" in metadata and metadata["task_type"] == "implement_feature": - features.append(interaction.get("content", "")) - - return "\n\n".join(features[:2]) # Return top 2 similar features - - async def get_agent_memory(self, agent_id: str, limit: int = 50) -> List[Dict]: - """Get agent's memory/interaction history""" - async with get_db_connection() as conn: - rows = await conn.fetch( - """ - SELECT content, metadata, created_at - FROM interactions - WHERE agent_id = $1 - ORDER BY created_at DESC - LIMIT $2 - """, - agent_id, - limit, - ) - return [dict(row) for row in rows] - - async def search_knowledge(self, query: str, limit: int = 10) -> List[Dict]: - """Search across all agent knowledge""" - return await self.get_similar_interactions(query, limit=limit) +import asyncio +import json +from typing import Dict, List, Optional + +import numpy as np +from sentence_transformers import SentenceTransformer + +from .database import get_db_connection + + +class ContextService: + def __init__(self): + # Load sentence transformer for embeddings + self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") + + def generate_embedding(self, text: str) -> List[float]: + """Generate embedding for text""" + embedding = self.embedding_model.encode(text) + return embedding.tolist() + + async def store_interaction( + self, agent_id: str, content: str, metadata: Dict = None + ) -> str: + """Store agent interaction with embedding""" + embedding = self.generate_embedding(content) + + async with get_db_connection() as conn: + interaction_id = await conn.fetchval( + """ + INSERT INTO interactions (agent_id, content, embedding, metadata) + VALUES ($1, $2, $3, $4) + RETURNING id + """, + agent_id, + content, + embedding, + metadata or {}, + ) + return str(interaction_id) + + async def get_similar_interactions( + self, + query: str, + agent_id: str = None, + limit: int = 5, + similarity_threshold: float = 0.7, + ) -> List[Dict]: + """Find similar interactions using vector similarity""" + query_embedding = self.generate_embedding(query) + + async with get_db_connection() as conn: + if agent_id: + rows = await conn.fetch( + """ + SELECT id, agent_id, content, metadata, created_at, + 1 - (embedding <=> $1) as similarity + FROM interactions + WHERE agent_id = $2 AND 1 - (embedding <=> $1) > $3 + ORDER BY similarity DESC + LIMIT $4 + """, + query_embedding, + agent_id, + similarity_threshold, + limit, + ) + else: + rows = await conn.fetch( + """ + SELECT id, agent_id, content, metadata, created_at, + 1 - (embedding <=> $1) as similarity + FROM interactions + WHERE 1 - (embedding <=> $1) > $2 + ORDER BY similarity DESC + LIMIT $3 + """, + query_embedding, + similarity_threshold, + limit, + ) + + return [dict(row) for row in rows] + + async def get_context(self, query: str, agent_id: str = None) -> Dict: + """Get relevant context for a query""" + similar_interactions = await self.get_similar_interactions(query, agent_id) + + # Get recent interactions from same agent + recent_interactions = [] + if agent_id: + async with get_db_connection() as conn: + rows = await conn.fetch( + """ + SELECT content, metadata, created_at + FROM interactions + WHERE agent_id = $1 + ORDER BY created_at DESC + LIMIT 10 + """, + agent_id, + ) + recent_interactions = [dict(row) for row in rows] + + return { + "similar_interactions": similar_interactions, + "recent_interactions": recent_interactions, + "relevant_code": self._extract_code_snippets(similar_interactions), + "similar_features": self._extract_similar_features(similar_interactions), + } + + def _extract_code_snippets(self, interactions: List[Dict]) -> str: + """Extract code snippets from interactions""" + code_snippets = [] + for interaction in interactions: + content = interaction.get("content", "") + # Simple extraction - look for code blocks + if "```" in content: + parts = content.split("```") + for i in range(1, len(parts), 2): + code_snippets.append(parts[i].strip()) + + return "\n\n".join(code_snippets[:3]) # Return top 3 snippets + + def _extract_similar_features(self, interactions: List[Dict]) -> str: + """Extract similar feature descriptions""" + features = [] + for interaction in interactions: + metadata = interaction.get("metadata", {}) + if "task_type" in metadata and metadata["task_type"] == "implement_feature": + features.append(interaction.get("content", "")) + + return "\n\n".join(features[:2]) # Return top 2 similar features + + async def get_agent_memory(self, agent_id: str, limit: int = 50) -> List[Dict]: + """Get agent's memory/interaction history""" + async with get_db_connection() as conn: + rows = await conn.fetch( + """ + SELECT content, metadata, created_at + FROM interactions + WHERE agent_id = $1 + ORDER BY created_at DESC + LIMIT $2 + """, + agent_id, + limit, + ) + return [dict(row) for row in rows] + + async def search_knowledge(self, query: str, limit: int = 10) -> List[Dict]: + """Search across all agent knowledge""" + return await self.get_similar_interactions(query, limit=limit) diff --git a/services/orchestrator/conversation_manager.py b/services/orchestrator/conversation_manager.py index e792c2c..507aa26 100644 --- a/services/orchestrator/conversation_manager.py +++ b/services/orchestrator/conversation_manager.py @@ -1,602 +1,602 @@ -""" -Conversation Manager for FuzeAgent Claude Code Integration - -Manages and stores complete conversations between agents and Claude Code, -providing comprehensive audit trails, debugging capabilities, and learning data. -""" - -import asyncio -import json -import logging -import time -import uuid -from dataclasses import dataclass -from datetime import datetime -from enum import Enum -from typing import Any, Dict, List, Optional - -# Importable both as `services.orchestrator.conversation_manager` (relative) and -# flat with services/orchestrator on sys.path (as the tests import it). -try: - from .database import get_db_connection -except ImportError: # pragma: no cover - flat import (no parent package) - from database import get_db_connection - -logger = logging.getLogger(__name__) - - -class MessageType(str, Enum): - USER_PROMPT = "user_prompt" - CLAUDE_RESPONSE = "claude_response" - SYSTEM_MESSAGE = "system_message" - ERROR_MESSAGE = "error_message" - CODE_EXECUTION = "code_execution" - TEST_RESULT = "test_result" - - -class InteractionType(str, Enum): - QUESTION = "question" - CLARIFICATION = "clarification" - APPROVAL_REQUEST = "approval_request" - ERROR_REPORT = "error_report" - PROGRESS_UPDATE = "progress_update" - - -@dataclass -class ConversationMessage: - """Represents a single message in a Claude Code conversation""" - - task_id: str - iteration_number: int - message_type: MessageType - content: str - token_count: Optional[int] = None - model_used: Optional[str] = None - temperature: Optional[float] = None - response_time_ms: Optional[int] = None - metadata: Optional[Dict[str, Any]] = None - - -@dataclass -class ConversationSession: - """Represents a complete conversation session""" - - agent_id: str - task_id: str - sandbox_id: str - session_started_at: datetime - session_ended_at: Optional[datetime] = None - total_messages: int = 0 - total_tokens: int = 0 - status: str = "active" - metadata: Optional[Dict[str, Any]] = None - - -class ConversationManager: - """ - Manages Claude Code conversations and provides comprehensive tracking. - - Features: - - Full conversation storage and retrieval - - Token usage tracking and cost analysis - - Human interaction management - - Code generation tracking - - Performance metrics collection - """ - - def __init__(self): - self.active_sessions: Dict[str, ConversationSession] = {} - - async def start_conversation_session( - self, - agent_id: str, - task_id: str, - sandbox_id: str, - metadata: Optional[Dict[str, Any]] = None, - ) -> str: - """Start a new conversation session for an agent""" - - session = ConversationSession( - agent_id=agent_id, - task_id=task_id, - sandbox_id=sandbox_id, - session_started_at=datetime.now(), - metadata=metadata or {}, - ) - - # Store session in database - session_id = await self._store_conversation_session(session) - session.metadata = session.metadata or {} - session.metadata["session_id"] = session_id - - # Track active session - self.active_sessions[session_id] = session - - logger.info( - f"Started conversation session {session_id} for agent {agent_id}, task {task_id}" - ) - return session_id - - async def end_conversation_session(self, session_id: str) -> bool: - """End a conversation session""" - - session = self.active_sessions.get(session_id) - if not session: - logger.warning(f"Session {session_id} not found in active sessions") - return False - - session.session_ended_at = datetime.now() - session.status = "completed" - - # Update database - await self._update_conversation_session(session_id, session) - - # Remove from active sessions - del self.active_sessions[session_id] - - logger.info(f"Ended conversation session {session_id}") - return True - - async def store_message( - self, - session_id: str, - message: ConversationMessage, - start_time: Optional[float] = None, - ) -> str: - """Store a conversation message""" - - # Calculate response time if start_time provided - if start_time and message.message_type == MessageType.CLAUDE_RESPONSE: - message.response_time_ms = int((time.time() - start_time) * 1000) - - # Store message in database - message_id = await self._store_claude_conversation(message) - - # Update session statistics - session = self.active_sessions.get(session_id) - if session: - session.total_messages += 1 - if message.token_count: - session.total_tokens += message.token_count - - logger.debug(f"Stored message {message_id} for session {session_id}") - return message_id - - async def store_user_prompt( - self, - session_id: str, - task_id: str, - iteration_number: int, - prompt: str, - model: str = "claude-3-5-sonnet-20241022", - temperature: float = 0.3, - metadata: Optional[Dict[str, Any]] = None, - ) -> str: - """Store a user prompt to Claude""" - - message = ConversationMessage( - task_id=task_id, - iteration_number=iteration_number, - message_type=MessageType.USER_PROMPT, - content=prompt, - model_used=model, - temperature=temperature, - metadata=metadata or {}, - ) - - return await self.store_message(session_id, message) - - async def store_claude_response( - self, - session_id: str, - task_id: str, - iteration_number: int, - response: str, - token_count: Optional[int] = None, - model: str = "claude-3-5-sonnet-20241022", - start_time: Optional[float] = None, - metadata: Optional[Dict[str, Any]] = None, - ) -> str: - """Store Claude's response""" - - message = ConversationMessage( - task_id=task_id, - iteration_number=iteration_number, - message_type=MessageType.CLAUDE_RESPONSE, - content=response, - token_count=token_count, - model_used=model, - metadata=metadata or {}, - ) - - return await self.store_message(session_id, message, start_time) - - async def store_code_generation( - self, - task_id: str, - iteration_number: int, - file_path: str, - file_type: str, - language: str, - content: str, - commit_hash: Optional[str] = None, - test_results: Optional[Dict[str, Any]] = None, - quality_metrics: Optional[Dict[str, Any]] = None, - ) -> str: - """Store generated code with metadata""" - - async with get_db_connection() as conn: - code_id = await conn.fetchval( - """ - INSERT INTO code_generations ( - task_id, iteration_number, file_path, file_type, language, - content, commit_hash, test_results, quality_metrics - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) - RETURNING id - """, - task_id, - iteration_number, - file_path, - file_type, - language, - content, - commit_hash, - json.dumps(test_results) if test_results else None, - json.dumps(quality_metrics) if quality_metrics else None, - ) - - logger.info(f"Stored code generation {code_id} for task {task_id}") - return str(code_id) - - async def store_human_interaction( - self, - task_id: str, - iteration_number: int, - interaction_type: InteractionType, - agent_message: str, - human_response: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, - ) -> str: - """Store human-in-the-loop interaction""" - - async with get_db_connection() as conn: - interaction_id = await conn.fetchval( - """ - INSERT INTO human_interactions ( - task_id, iteration_number, interaction_type, - agent_message, human_response, metadata - ) VALUES ($1, $2, $3, $4, $5, $6) - RETURNING id - """, - task_id, - iteration_number, - interaction_type.value, - agent_message, - human_response, - json.dumps(metadata) if metadata else {}, - ) - - logger.info(f"Stored human interaction {interaction_id} for task {task_id}") - return str(interaction_id) - - async def update_human_response( - self, interaction_id: str, human_response: str - ) -> bool: - """Update human response to an interaction""" - - async with get_db_connection() as conn: - result = await conn.execute( - """ - UPDATE human_interactions - SET human_response = $1, - responded_at = CURRENT_TIMESTAMP, - response_time_seconds = EXTRACT(EPOCH FROM (CURRENT_TIMESTAMP - asked_at)) - WHERE id = $2 - """, - human_response, - interaction_id, - ) - - success = result != "UPDATE 0" - if success: - logger.info(f"Updated human response for interaction {interaction_id}") - return success - - async def store_performance_metric( - self, - agent_id: str, - task_id: str, - metric_type: str, - metric_value: float, - metric_unit: Optional[str] = None, - context: Optional[Dict[str, Any]] = None, - ) -> str: - """Store agent performance metric""" - - async with get_db_connection() as conn: - metric_id = await conn.fetchval( - """ - INSERT INTO agent_performance_metrics ( - agent_id, task_id, metric_type, metric_value, - metric_unit, context - ) VALUES ($1, $2, $3, $4, $5, $6) - RETURNING id - """, - agent_id, - task_id, - metric_type, - metric_value, - metric_unit, - json.dumps(context) if context else {}, - ) - - logger.debug( - f"Stored performance metric {metric_id}: {metric_type}={metric_value}" - ) - return str(metric_id) - - async def get_conversation_history( - self, - task_id: str, - iteration_number: Optional[int] = None, - message_types: Optional[List[MessageType]] = None, - limit: Optional[int] = None, - ) -> List[Dict[str, Any]]: - """Get conversation history for a task""" - - conditions = ["task_id = $1"] - params = [task_id] - param_count = 1 - - if iteration_number is not None: - param_count += 1 - conditions.append(f"iteration_number = ${param_count}") - params.append(iteration_number) - - if message_types: - param_count += 1 - conditions.append(f"message_type = ANY(${param_count})") - params.append([mt.value for mt in message_types]) - - where_clause = " AND ".join(conditions) - limit_clause = "" - if limit: - param_count += 1 - limit_clause = f"LIMIT ${param_count}" - params.append(limit) - - async with get_db_connection() as conn: - rows = await conn.fetch( - f""" - SELECT * FROM claude_conversations - WHERE {where_clause} - ORDER BY created_at ASC - {limit_clause} - """, # nosec B608 -- where/limit clauses are fixed fragments with $N placeholders; all values bound as query params - *params, - ) - - return [dict(row) for row in rows] - - async def get_conversation_summary(self, task_id: str) -> Dict[str, Any]: - """Get conversation summary with statistics""" - - async with get_db_connection() as conn: - # Get message statistics - stats = await conn.fetchrow( - """ - SELECT - COUNT(*) as total_messages, - SUM(token_count) as total_tokens, - AVG(response_time_ms) as avg_response_time, - MAX(iteration_number) as max_iteration - FROM claude_conversations - WHERE task_id = $1 - """, - task_id, - ) - - # Get message type breakdown - type_breakdown = await conn.fetch( - """ - SELECT message_type, COUNT(*) as count - FROM claude_conversations - WHERE task_id = $1 - GROUP BY message_type - """, - task_id, - ) - - # Get human interactions - human_interactions = await conn.fetchrow( - """ - SELECT - COUNT(*) as total_interactions, - COUNT(human_response) as responded_interactions, - AVG(response_time_seconds) as avg_response_time - FROM human_interactions - WHERE task_id = $1 - """, - task_id, - ) - - return { - "task_id": task_id, - "total_messages": stats["total_messages"] or 0, - "total_tokens": stats["total_tokens"] or 0, - "avg_response_time_ms": float(stats["avg_response_time"] or 0), - "max_iteration": stats["max_iteration"] or 0, - "message_types": { - row["message_type"]: row["count"] for row in type_breakdown - }, - "human_interactions": { - "total": human_interactions["total_interactions"] or 0, - "responded": human_interactions["responded_interactions"] or 0, - "avg_response_time_seconds": float( - human_interactions["avg_response_time"] or 0 - ), - }, - } - - async def get_code_generations( - self, - task_id: str, - iteration_number: Optional[int] = None, - file_type: Optional[str] = None, - language: Optional[str] = None, - ) -> List[Dict[str, Any]]: - """Get code generations for a task""" - - conditions = ["task_id = $1"] - params = [task_id] - param_count = 1 - - if iteration_number is not None: - param_count += 1 - conditions.append(f"iteration_number = ${param_count}") - params.append(iteration_number) - - if file_type: - param_count += 1 - conditions.append(f"file_type = ${param_count}") - params.append(file_type) - - if language: - param_count += 1 - conditions.append(f"language = ${param_count}") - params.append(language) - - where_clause = " AND ".join(conditions) - - async with get_db_connection() as conn: - rows = await conn.fetch( - f""" - SELECT * FROM code_generations - WHERE {where_clause} - ORDER BY generated_at ASC - """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params - *params, - ) - - return [dict(row) for row in rows] - - async def get_agent_performance_metrics( - self, - agent_id: Optional[str] = None, - task_id: Optional[str] = None, - metric_types: Optional[List[str]] = None, - time_range_hours: Optional[int] = None, - ) -> List[Dict[str, Any]]: - """Get agent performance metrics""" - - conditions = [] - params = [] - param_count = 0 - - if agent_id: - param_count += 1 - conditions.append(f"agent_id = ${param_count}") - params.append(agent_id) - - if task_id: - param_count += 1 - conditions.append(f"task_id = ${param_count}") - params.append(task_id) - - if metric_types: - param_count += 1 - conditions.append(f"metric_type = ANY(${param_count})") - params.append(metric_types) - - if time_range_hours: - param_count += 1 - conditions.append( - f"measured_at >= NOW() - (INTERVAL '1 hour' * ${param_count})" - ) - params.append(time_range_hours) - - where_clause = "WHERE " + " AND ".join(conditions) if conditions else "" - - async with get_db_connection() as conn: - rows = await conn.fetch( - f""" - SELECT * FROM agent_performance_metrics - {where_clause} - ORDER BY measured_at DESC - """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values (incl. interval multiplier) bound as query params - *params, - ) - - return [dict(row) for row in rows] - - # Private methods - - async def _store_conversation_session(self, session: ConversationSession) -> str: - """Store conversation session in database""" - - async with get_db_connection() as conn: - session_id = await conn.fetchval( - """ - INSERT INTO agent_conversation_sessions ( - agent_id, task_id, sandbox_id, session_started_at, - total_messages, total_tokens, status, metadata - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) - RETURNING id - """, - session.agent_id, - session.task_id, - session.sandbox_id, - session.session_started_at, - session.total_messages, - session.total_tokens, - session.status, - json.dumps(session.metadata) if session.metadata else {}, - ) - - return str(session_id) - - async def _update_conversation_session( - self, session_id: str, session: ConversationSession - ): - """Update conversation session in database""" - - async with get_db_connection() as conn: - await conn.execute( - """ - UPDATE agent_conversation_sessions - SET session_ended_at = $1, total_messages = $2, - total_tokens = $3, status = $4, metadata = $5 - WHERE id = $6 - """, - session.session_ended_at, - session.total_messages, - session.total_tokens, - session.status, - json.dumps(session.metadata) if session.metadata else {}, - session_id, - ) - - async def _store_claude_conversation(self, message: ConversationMessage) -> str: - """Store Claude conversation message in database""" - - async with get_db_connection() as conn: - message_id = await conn.fetchval( - """ - INSERT INTO claude_conversations ( - task_id, iteration_number, message_type, content, - token_count, model_used, temperature, response_time_ms, metadata - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) - RETURNING id - """, - message.task_id, - message.iteration_number, - message.message_type.value, - message.content, - message.token_count, - message.model_used, - message.temperature, - message.response_time_ms, - json.dumps(message.metadata) if message.metadata else {}, - ) - - return str(message_id) +""" +Conversation Manager for FuzeAgent Claude Code Integration + +Manages and stores complete conversations between agents and Claude Code, +providing comprehensive audit trails, debugging capabilities, and learning data. +""" + +import asyncio +import json +import logging +import time +import uuid +from dataclasses import dataclass +from datetime import datetime +from enum import Enum +from typing import Any, Dict, List, Optional + +# Importable both as `services.orchestrator.conversation_manager` (relative) and +# flat with services/orchestrator on sys.path (as the tests import it). +try: + from .database import get_db_connection +except ImportError: # pragma: no cover - flat import (no parent package) + from database import get_db_connection + +logger = logging.getLogger(__name__) + + +class MessageType(str, Enum): + USER_PROMPT = "user_prompt" + CLAUDE_RESPONSE = "claude_response" + SYSTEM_MESSAGE = "system_message" + ERROR_MESSAGE = "error_message" + CODE_EXECUTION = "code_execution" + TEST_RESULT = "test_result" + + +class InteractionType(str, Enum): + QUESTION = "question" + CLARIFICATION = "clarification" + APPROVAL_REQUEST = "approval_request" + ERROR_REPORT = "error_report" + PROGRESS_UPDATE = "progress_update" + + +@dataclass +class ConversationMessage: + """Represents a single message in a Claude Code conversation""" + + task_id: str + iteration_number: int + message_type: MessageType + content: str + token_count: Optional[int] = None + model_used: Optional[str] = None + temperature: Optional[float] = None + response_time_ms: Optional[int] = None + metadata: Optional[Dict[str, Any]] = None + + +@dataclass +class ConversationSession: + """Represents a complete conversation session""" + + agent_id: str + task_id: str + sandbox_id: str + session_started_at: datetime + session_ended_at: Optional[datetime] = None + total_messages: int = 0 + total_tokens: int = 0 + status: str = "active" + metadata: Optional[Dict[str, Any]] = None + + +class ConversationManager: + """ + Manages Claude Code conversations and provides comprehensive tracking. + + Features: + - Full conversation storage and retrieval + - Token usage tracking and cost analysis + - Human interaction management + - Code generation tracking + - Performance metrics collection + """ + + def __init__(self): + self.active_sessions: Dict[str, ConversationSession] = {} + + async def start_conversation_session( + self, + agent_id: str, + task_id: str, + sandbox_id: str, + metadata: Optional[Dict[str, Any]] = None, + ) -> str: + """Start a new conversation session for an agent""" + + session = ConversationSession( + agent_id=agent_id, + task_id=task_id, + sandbox_id=sandbox_id, + session_started_at=datetime.now(), + metadata=metadata or {}, + ) + + # Store session in database + session_id = await self._store_conversation_session(session) + session.metadata = session.metadata or {} + session.metadata["session_id"] = session_id + + # Track active session + self.active_sessions[session_id] = session + + logger.info( + f"Started conversation session {session_id} for agent {agent_id}, task {task_id}" + ) + return session_id + + async def end_conversation_session(self, session_id: str) -> bool: + """End a conversation session""" + + session = self.active_sessions.get(session_id) + if not session: + logger.warning(f"Session {session_id} not found in active sessions") + return False + + session.session_ended_at = datetime.now() + session.status = "completed" + + # Update database + await self._update_conversation_session(session_id, session) + + # Remove from active sessions + del self.active_sessions[session_id] + + logger.info(f"Ended conversation session {session_id}") + return True + + async def store_message( + self, + session_id: str, + message: ConversationMessage, + start_time: Optional[float] = None, + ) -> str: + """Store a conversation message""" + + # Calculate response time if start_time provided + if start_time and message.message_type == MessageType.CLAUDE_RESPONSE: + message.response_time_ms = int((time.time() - start_time) * 1000) + + # Store message in database + message_id = await self._store_claude_conversation(message) + + # Update session statistics + session = self.active_sessions.get(session_id) + if session: + session.total_messages += 1 + if message.token_count: + session.total_tokens += message.token_count + + logger.debug(f"Stored message {message_id} for session {session_id}") + return message_id + + async def store_user_prompt( + self, + session_id: str, + task_id: str, + iteration_number: int, + prompt: str, + model: str = "claude-3-5-sonnet-20241022", + temperature: float = 0.3, + metadata: Optional[Dict[str, Any]] = None, + ) -> str: + """Store a user prompt to Claude""" + + message = ConversationMessage( + task_id=task_id, + iteration_number=iteration_number, + message_type=MessageType.USER_PROMPT, + content=prompt, + model_used=model, + temperature=temperature, + metadata=metadata or {}, + ) + + return await self.store_message(session_id, message) + + async def store_claude_response( + self, + session_id: str, + task_id: str, + iteration_number: int, + response: str, + token_count: Optional[int] = None, + model: str = "claude-3-5-sonnet-20241022", + start_time: Optional[float] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> str: + """Store Claude's response""" + + message = ConversationMessage( + task_id=task_id, + iteration_number=iteration_number, + message_type=MessageType.CLAUDE_RESPONSE, + content=response, + token_count=token_count, + model_used=model, + metadata=metadata or {}, + ) + + return await self.store_message(session_id, message, start_time) + + async def store_code_generation( + self, + task_id: str, + iteration_number: int, + file_path: str, + file_type: str, + language: str, + content: str, + commit_hash: Optional[str] = None, + test_results: Optional[Dict[str, Any]] = None, + quality_metrics: Optional[Dict[str, Any]] = None, + ) -> str: + """Store generated code with metadata""" + + async with get_db_connection() as conn: + code_id = await conn.fetchval( + """ + INSERT INTO code_generations ( + task_id, iteration_number, file_path, file_type, language, + content, commit_hash, test_results, quality_metrics + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) + RETURNING id + """, + task_id, + iteration_number, + file_path, + file_type, + language, + content, + commit_hash, + json.dumps(test_results) if test_results else None, + json.dumps(quality_metrics) if quality_metrics else None, + ) + + logger.info(f"Stored code generation {code_id} for task {task_id}") + return str(code_id) + + async def store_human_interaction( + self, + task_id: str, + iteration_number: int, + interaction_type: InteractionType, + agent_message: str, + human_response: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> str: + """Store human-in-the-loop interaction""" + + async with get_db_connection() as conn: + interaction_id = await conn.fetchval( + """ + INSERT INTO human_interactions ( + task_id, iteration_number, interaction_type, + agent_message, human_response, metadata + ) VALUES ($1, $2, $3, $4, $5, $6) + RETURNING id + """, + task_id, + iteration_number, + interaction_type.value, + agent_message, + human_response, + json.dumps(metadata) if metadata else {}, + ) + + logger.info(f"Stored human interaction {interaction_id} for task {task_id}") + return str(interaction_id) + + async def update_human_response( + self, interaction_id: str, human_response: str + ) -> bool: + """Update human response to an interaction""" + + async with get_db_connection() as conn: + result = await conn.execute( + """ + UPDATE human_interactions + SET human_response = $1, + responded_at = CURRENT_TIMESTAMP, + response_time_seconds = EXTRACT(EPOCH FROM (CURRENT_TIMESTAMP - asked_at)) + WHERE id = $2 + """, + human_response, + interaction_id, + ) + + success = result != "UPDATE 0" + if success: + logger.info(f"Updated human response for interaction {interaction_id}") + return success + + async def store_performance_metric( + self, + agent_id: str, + task_id: str, + metric_type: str, + metric_value: float, + metric_unit: Optional[str] = None, + context: Optional[Dict[str, Any]] = None, + ) -> str: + """Store agent performance metric""" + + async with get_db_connection() as conn: + metric_id = await conn.fetchval( + """ + INSERT INTO agent_performance_metrics ( + agent_id, task_id, metric_type, metric_value, + metric_unit, context + ) VALUES ($1, $2, $3, $4, $5, $6) + RETURNING id + """, + agent_id, + task_id, + metric_type, + metric_value, + metric_unit, + json.dumps(context) if context else {}, + ) + + logger.debug( + f"Stored performance metric {metric_id}: {metric_type}={metric_value}" + ) + return str(metric_id) + + async def get_conversation_history( + self, + task_id: str, + iteration_number: Optional[int] = None, + message_types: Optional[List[MessageType]] = None, + limit: Optional[int] = None, + ) -> List[Dict[str, Any]]: + """Get conversation history for a task""" + + conditions = ["task_id = $1"] + params = [task_id] + param_count = 1 + + if iteration_number is not None: + param_count += 1 + conditions.append(f"iteration_number = ${param_count}") + params.append(iteration_number) + + if message_types: + param_count += 1 + conditions.append(f"message_type = ANY(${param_count})") + params.append([mt.value for mt in message_types]) + + where_clause = " AND ".join(conditions) + limit_clause = "" + if limit: + param_count += 1 + limit_clause = f"LIMIT ${param_count}" + params.append(limit) + + async with get_db_connection() as conn: + rows = await conn.fetch( + f""" + SELECT * FROM claude_conversations + WHERE {where_clause} + ORDER BY created_at ASC + {limit_clause} + """, # nosec B608 -- where/limit clauses are fixed fragments with $N placeholders; all values bound as query params + *params, + ) + + return [dict(row) for row in rows] + + async def get_conversation_summary(self, task_id: str) -> Dict[str, Any]: + """Get conversation summary with statistics""" + + async with get_db_connection() as conn: + # Get message statistics + stats = await conn.fetchrow( + """ + SELECT + COUNT(*) as total_messages, + SUM(token_count) as total_tokens, + AVG(response_time_ms) as avg_response_time, + MAX(iteration_number) as max_iteration + FROM claude_conversations + WHERE task_id = $1 + """, + task_id, + ) + + # Get message type breakdown + type_breakdown = await conn.fetch( + """ + SELECT message_type, COUNT(*) as count + FROM claude_conversations + WHERE task_id = $1 + GROUP BY message_type + """, + task_id, + ) + + # Get human interactions + human_interactions = await conn.fetchrow( + """ + SELECT + COUNT(*) as total_interactions, + COUNT(human_response) as responded_interactions, + AVG(response_time_seconds) as avg_response_time + FROM human_interactions + WHERE task_id = $1 + """, + task_id, + ) + + return { + "task_id": task_id, + "total_messages": stats["total_messages"] or 0, + "total_tokens": stats["total_tokens"] or 0, + "avg_response_time_ms": float(stats["avg_response_time"] or 0), + "max_iteration": stats["max_iteration"] or 0, + "message_types": { + row["message_type"]: row["count"] for row in type_breakdown + }, + "human_interactions": { + "total": human_interactions["total_interactions"] or 0, + "responded": human_interactions["responded_interactions"] or 0, + "avg_response_time_seconds": float( + human_interactions["avg_response_time"] or 0 + ), + }, + } + + async def get_code_generations( + self, + task_id: str, + iteration_number: Optional[int] = None, + file_type: Optional[str] = None, + language: Optional[str] = None, + ) -> List[Dict[str, Any]]: + """Get code generations for a task""" + + conditions = ["task_id = $1"] + params = [task_id] + param_count = 1 + + if iteration_number is not None: + param_count += 1 + conditions.append(f"iteration_number = ${param_count}") + params.append(iteration_number) + + if file_type: + param_count += 1 + conditions.append(f"file_type = ${param_count}") + params.append(file_type) + + if language: + param_count += 1 + conditions.append(f"language = ${param_count}") + params.append(language) + + where_clause = " AND ".join(conditions) + + async with get_db_connection() as conn: + rows = await conn.fetch( + f""" + SELECT * FROM code_generations + WHERE {where_clause} + ORDER BY generated_at ASC + """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params + *params, + ) + + return [dict(row) for row in rows] + + async def get_agent_performance_metrics( + self, + agent_id: Optional[str] = None, + task_id: Optional[str] = None, + metric_types: Optional[List[str]] = None, + time_range_hours: Optional[int] = None, + ) -> List[Dict[str, Any]]: + """Get agent performance metrics""" + + conditions = [] + params = [] + param_count = 0 + + if agent_id: + param_count += 1 + conditions.append(f"agent_id = ${param_count}") + params.append(agent_id) + + if task_id: + param_count += 1 + conditions.append(f"task_id = ${param_count}") + params.append(task_id) + + if metric_types: + param_count += 1 + conditions.append(f"metric_type = ANY(${param_count})") + params.append(metric_types) + + if time_range_hours: + param_count += 1 + conditions.append( + f"measured_at >= NOW() - (INTERVAL '1 hour' * ${param_count})" + ) + params.append(time_range_hours) + + where_clause = "WHERE " + " AND ".join(conditions) if conditions else "" + + async with get_db_connection() as conn: + rows = await conn.fetch( + f""" + SELECT * FROM agent_performance_metrics + {where_clause} + ORDER BY measured_at DESC + """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values (incl. interval multiplier) bound as query params + *params, + ) + + return [dict(row) for row in rows] + + # Private methods + + async def _store_conversation_session(self, session: ConversationSession) -> str: + """Store conversation session in database""" + + async with get_db_connection() as conn: + session_id = await conn.fetchval( + """ + INSERT INTO agent_conversation_sessions ( + agent_id, task_id, sandbox_id, session_started_at, + total_messages, total_tokens, status, metadata + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + RETURNING id + """, + session.agent_id, + session.task_id, + session.sandbox_id, + session.session_started_at, + session.total_messages, + session.total_tokens, + session.status, + json.dumps(session.metadata) if session.metadata else {}, + ) + + return str(session_id) + + async def _update_conversation_session( + self, session_id: str, session: ConversationSession + ): + """Update conversation session in database""" + + async with get_db_connection() as conn: + await conn.execute( + """ + UPDATE agent_conversation_sessions + SET session_ended_at = $1, total_messages = $2, + total_tokens = $3, status = $4, metadata = $5 + WHERE id = $6 + """, + session.session_ended_at, + session.total_messages, + session.total_tokens, + session.status, + json.dumps(session.metadata) if session.metadata else {}, + session_id, + ) + + async def _store_claude_conversation(self, message: ConversationMessage) -> str: + """Store Claude conversation message in database""" + + async with get_db_connection() as conn: + message_id = await conn.fetchval( + """ + INSERT INTO claude_conversations ( + task_id, iteration_number, message_type, content, + token_count, model_used, temperature, response_time_ms, metadata + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) + RETURNING id + """, + message.task_id, + message.iteration_number, + message.message_type.value, + message.content, + message.token_count, + message.model_used, + message.temperature, + message.response_time_ms, + json.dumps(message.metadata) if message.metadata else {}, + ) + + return str(message_id) diff --git a/services/orchestrator/coordination_endpoints.py b/services/orchestrator/coordination_endpoints.py index cdcd244..ee08b34 100644 --- a/services/orchestrator/coordination_endpoints.py +++ b/services/orchestrator/coordination_endpoints.py @@ -1,625 +1,625 @@ -""" -Cross-Product Coordination API Endpoints - -This module provides REST API endpoints for managing cross-product coordination -within the WCG ecosystem. Enables centralized coordination between FuzeAgent, -FuzeFront, HubHit, DeployAI, and other WCG products. -""" - -import logging -from datetime import date, datetime -from typing import Any, Dict, List, Optional - -import asyncpg -from fastapi import APIRouter, Depends, HTTPException, Path, Query -from pydantic import BaseModel, Field - -from .database import get_db_connection - -logger = logging.getLogger(__name__) -router = APIRouter(prefix="/coordination", tags=["Cross-Product Coordination"]) - - -# Pydantic Models -class ProductRegistration(BaseModel): - id: str = Field(..., description="Unique product identifier") - name: str = Field(..., description="Product display name") - version: str = Field(..., description="Current product version") - endpoints: List[str] = Field(default=[], description="API endpoints exposed") - dependencies: List[str] = Field(default=[], description="Product dependencies") - resource_requirements: Dict[str, Any] = Field( - default={}, description="Resource needs" - ) - team_contacts: List[str] = Field(default=[], description="Team contact information") - priority_level: int = Field(default=5, ge=1, le=10, description="Business priority") - metadata: Dict[str, Any] = Field(default={}, description="Additional metadata") - - -class CoordinationRequestCreate(BaseModel): - requesting_product: str = Field(..., description="Product making the request") - target_products: List[str] = Field( - ..., description="Target products for coordination" - ) - coordination_type: str = Field(..., description="Type of coordination needed") - scope: str = Field(default="product_group", description="Coordination scope") - priority: str = Field(default="medium", description="Request priority") - title: str = Field(..., description="Brief title for the request") - description: str = Field(..., description="Detailed description") - resource_requirements: Dict[str, Any] = Field( - default={}, description="Required resources" - ) - proposed_timeline: Dict[str, Any] = Field( - default={}, description="Proposed timeline" - ) - - -class CoordinationResolution(BaseModel): - resolution_plan: Dict[str, Any] = Field(..., description="Detailed resolution plan") - resolver_id: str = Field(..., description="ID of the resolver") - notes: Optional[str] = Field(None, description="Additional resolution notes") - - -class ResourceAllocationCreate(BaseModel): - product_id: str = Field(..., description="Product requesting allocation") - resource_type: str = Field(..., description="Type of resource") - resource_name: str = Field(..., description="Specific resource name") - allocation_details: Dict[str, Any] = Field( - default={}, description="Allocation specifics" - ) - valid_until: Optional[datetime] = Field(None, description="Allocation expiry") - - -# Product Registration Endpoints -@router.post( - "/products/register", summary="Register a new product in coordination system" -) -async def register_product(product: ProductRegistration): - """Register a new product in the cross-product coordination system""" - try: - async with get_db_connection() as conn: - await conn.execute( - """ - INSERT INTO product_registry ( - id, name, version, endpoints, dependencies, - resource_requirements, team_contacts, priority_level, - metadata, registered_at, updated_at - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) - ON CONFLICT (id) DO UPDATE SET - name = $2, version = $3, endpoints = $4, dependencies = $5, - resource_requirements = $6, team_contacts = $7, - priority_level = $8, metadata = $9, updated_at = $11 - """, - product.id, - product.name, - product.version, - product.endpoints, - product.dependencies, - product.resource_requirements, - product.team_contacts, - product.priority_level, - product.metadata, - datetime.utcnow(), - datetime.utcnow(), - ) - - return { - "status": "success", - "product_id": product.id, - "message": "Product registered successfully", - } - - except Exception as e: - logger.error(f"Error registering product {product.id}: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to register product: {str(e)}" - ) - - -@router.get("/products", summary="List all registered products") -async def list_products( - priority_min: int = Query(1, ge=1, le=10, description="Minimum priority level"), - active_only: bool = Query(True, description="Show only active products"), -): - """List all products registered in the coordination system""" - try: - async with get_db_connection() as conn: - products = await conn.fetch( - """ - SELECT id, name, version, priority_level, metadata, - endpoints, dependencies, registered_at, updated_at - FROM product_registry - WHERE priority_level >= $1 - ORDER BY priority_level DESC, name ASC - """, - priority_min, - ) - - return { - "products": [dict(product) for product in products], - "total_count": len(products), - } - - except Exception as e: - logger.error(f"Error listing products: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to list products: {str(e)}" - ) - - -@router.get("/products/{product_id}", summary="Get specific product details") -async def get_product(product_id: str = Path(..., description="Product ID")): - """Get detailed information about a specific product""" - try: - async with get_db_connection() as conn: - product = await conn.fetchrow( - """ - SELECT * FROM product_registry WHERE id = $1 - """, - product_id, - ) - - if not product: - raise HTTPException( - status_code=404, detail=f"Product {product_id} not found" - ) - - # Get active coordination requests involving this product - async with get_db_connection() as conn: - coordination_requests = await conn.fetch( - """ - SELECT id, title, coordination_type, priority, status, created_at - FROM coordination_requests - WHERE requesting_product = $1 - OR $1 = ANY(string_to_array(replace(replace(target_products::text, '[', ''), ']', ''), ',')) - ORDER BY created_at DESC LIMIT 10 - """, - product_id, - ) - - return { - "product": dict(product), - "active_coordination_requests": [ - dict(req) for req in coordination_requests - ], - } - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting product {product_id}: {e}") - raise HTTPException(status_code=500, detail=f"Failed to get product: {str(e)}") - - -# Coordination Request Endpoints -@router.post("/requests", summary="Create a new coordination request") -async def create_coordination_request(request: CoordinationRequestCreate): - """Create a new cross-product coordination request""" - try: - # Validate requesting product exists - async with get_db_connection() as conn: - requesting_product = await conn.fetchrow( - """ - SELECT id FROM product_registry WHERE id = $1 - """, - request.requesting_product, - ) - - if not requesting_product: - raise HTTPException( - status_code=400, - detail=f"Requesting product {request.requesting_product} not found", - ) - - # Create coordination request - async with get_db_connection() as conn: - request_id = await conn.fetchval( - """ - INSERT INTO coordination_requests ( - requesting_product, target_products, coordination_type, - scope, priority, title, description, resource_requirements, - proposed_timeline, stakeholders, created_at, updated_at - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) - RETURNING id - """, - request.requesting_product, - request.target_products, - request.coordination_type, - request.scope, - request.priority, - request.title, - request.description, - request.resource_requirements, - request.proposed_timeline, - [], # stakeholders - could be auto-populated - datetime.utcnow(), - datetime.utcnow(), - ) - - # Log coordination history - async with get_db_connection() as conn: - await conn.execute( - """ - INSERT INTO coordination_history ( - coordination_request_id, action, actor_type, details - ) VALUES ($1, $2, $3, $4) - """, - request_id, - "created", - "system", - {"created_via": "api"}, - ) - - return { - "status": "success", - "request_id": str(request_id), - "message": "Coordination request created successfully", - } - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error creating coordination request: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to create coordination request: {str(e)}" - ) - - -@router.get("/requests", summary="List coordination requests") -async def list_coordination_requests( - status: Optional[str] = Query(None, description="Filter by status"), - priority: Optional[str] = Query(None, description="Filter by priority"), - product_id: Optional[str] = Query(None, description="Filter by product"), - limit: int = Query(50, ge=1, le=200, description="Max number of results"), -): - """List coordination requests with optional filters""" - try: - where_conditions = [] - params = [] - param_count = 0 - - if status: - param_count += 1 - where_conditions.append(f"status = ${param_count}") - params.append(status) - - if priority: - param_count += 1 - where_conditions.append(f"priority = ${param_count}") - params.append(priority) - - if product_id: - param_count += 1 - where_conditions.append( - f"(requesting_product = ${param_count} OR ${param_count} = ANY(string_to_array(replace(replace(target_products::text, '[', ''), ']', ''), ',')))" - ) - params.append(product_id) - - where_clause = ( - " WHERE " + " AND ".join(where_conditions) if where_conditions else "" - ) - param_count += 1 - params.append(limit) - - query = f""" - SELECT id, requesting_product, target_products, coordination_type, - scope, priority, status, title, description, created_at, updated_at - FROM coordination_requests - {where_clause} - ORDER BY - CASE priority - WHEN 'critical' THEN 1 - WHEN 'high' THEN 2 - WHEN 'medium' THEN 3 - WHEN 'low' THEN 4 - END, - created_at DESC - LIMIT ${param_count} - """ # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params - - async with get_db_connection() as conn: - requests = await conn.fetch(query, *params) - - return { - "coordination_requests": [dict(req) for req in requests], - "total_count": len(requests), - } - - except Exception as e: - logger.error(f"Error listing coordination requests: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to list coordination requests: {str(e)}" - ) - - -@router.get("/requests/{request_id}", summary="Get coordination request details") -async def get_coordination_request( - request_id: str = Path(..., description="Coordination request ID") -): - """Get detailed information about a specific coordination request""" - try: - async with get_db_connection() as conn: - request = await conn.fetchrow( - """ - SELECT * FROM coordination_requests WHERE id = $1 - """, - request_id, - ) - - if not request: - raise HTTPException( - status_code=404, detail=f"Coordination request {request_id} not found" - ) - - # Get coordination history - async with get_db_connection() as conn: - history = await conn.fetch( - """ - SELECT action, actor_id, actor_type, details, timestamp - FROM coordination_history - WHERE coordination_request_id = $1 - ORDER BY timestamp ASC - """, - request_id, - ) - - return { - "coordination_request": dict(request), - "history": [dict(h) for h in history], - } - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting coordination request {request_id}: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to get coordination request: {str(e)}" - ) - - -@router.put("/requests/{request_id}/resolve", summary="Resolve coordination request") -async def resolve_coordination_request( - request_id: str = Path(..., description="Coordination request ID"), - resolution: CoordinationResolution = ..., -): - """Resolve a coordination request with a specific plan""" - try: - async with get_db_connection() as conn: - # Check if request exists and is pending - existing_request = await conn.fetchrow( - """ - SELECT id, status FROM coordination_requests WHERE id = $1 - """, - request_id, - ) - - if not existing_request: - raise HTTPException( - status_code=404, detail=f"Coordination request {request_id} not found" - ) - - if existing_request["status"] not in ["pending", "in_progress"]: - raise HTTPException( - status_code=400, - detail=f"Request is already {existing_request['status']}", - ) - - # Update request status and resolution - async with get_db_connection() as conn: - await conn.execute( - """ - UPDATE coordination_requests SET - status = 'resolved', - resolution_plan = $2, - resolved_at = $3, - updated_at = $4 - WHERE id = $1 - """, - request_id, - resolution.resolution_plan, - datetime.utcnow(), - datetime.utcnow(), - ) - - # Log resolution in history - async with get_db_connection() as conn: - await conn.execute( - """ - INSERT INTO coordination_history ( - coordination_request_id, action, actor_id, actor_type, details - ) VALUES ($1, $2, $3, $4, $5) - """, - request_id, - "resolved", - resolution.resolver_id, - "agent", - { - "resolution_plan": resolution.resolution_plan, - "notes": resolution.notes, - }, - ) - - return { - "status": "success", - "request_id": request_id, - "message": "Coordination request resolved successfully", - } - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error resolving coordination request {request_id}: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to resolve coordination request: {str(e)}" - ) - - -# Resource Management Endpoints -@router.post("/resources/allocate", summary="Allocate resources to a product") -async def allocate_resource(allocation: ResourceAllocationCreate): - """Allocate a resource to a specific product""" - try: - async with get_db_connection() as conn: - allocation_id = await conn.fetchval( - """ - INSERT INTO resource_allocations ( - product_id, resource_type, resource_name, allocation_details, - valid_until, status, created_at, updated_at - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) - RETURNING id - """, - allocation.product_id, - allocation.resource_type, - allocation.resource_name, - allocation.allocation_details, - allocation.valid_until, - "active", - datetime.utcnow(), - datetime.utcnow(), - ) - - return { - "status": "success", - "allocation_id": str(allocation_id), - "message": "Resource allocated successfully", - } - - except asyncpg.UniqueViolationError: - raise HTTPException( - status_code=409, - detail=f"Resource {allocation.resource_name} of type {allocation.resource_type} is already allocated", - ) - except Exception as e: - logger.error(f"Error allocating resource: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to allocate resource: {str(e)}" - ) - - -@router.get("/resources", summary="List resource allocations") -async def list_resource_allocations( - product_id: Optional[str] = Query(None, description="Filter by product"), - resource_type: Optional[str] = Query(None, description="Filter by resource type"), - status: Optional[str] = Query(None, description="Filter by status"), -): - """List current resource allocations""" - try: - where_conditions = [] - params = [] - param_count = 0 - - if product_id: - param_count += 1 - where_conditions.append(f"product_id = ${param_count}") - params.append(product_id) - - if resource_type: - param_count += 1 - where_conditions.append(f"resource_type = ${param_count}") - params.append(resource_type) - - if status: - param_count += 1 - where_conditions.append(f"status = ${param_count}") - params.append(status) - - where_clause = ( - " WHERE " + " AND ".join(where_conditions) if where_conditions else "" - ) - - query = f""" - SELECT ra.*, pr.name as product_name - FROM resource_allocations ra - LEFT JOIN product_registry pr ON ra.product_id = pr.id - {where_clause} - ORDER BY ra.created_at DESC - """ # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params - - async with get_db_connection() as conn: - allocations = await conn.fetch(query, *params) - - return { - "resource_allocations": [dict(alloc) for alloc in allocations], - "total_count": len(allocations), - } - - except Exception as e: - logger.error(f"Error listing resource allocations: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to list resource allocations: {str(e)}" - ) - - -# Dashboard and Status Endpoints -@router.get("/status", summary="Get overall coordination system status") -async def get_coordination_status(): - """Get overall status of the cross-product coordination system""" - try: - async with get_db_connection() as conn: - # Get request counts by status - request_stats = await conn.fetch(""" - SELECT status, priority, COUNT(*) as count - FROM coordination_requests - GROUP BY status, priority - ORDER BY status, priority - """) - - # Get product count - product_count = await conn.fetchval(""" - SELECT COUNT(*) FROM product_registry - """) - - # Get resource allocation stats - resource_stats = await conn.fetch(""" - SELECT resource_type, status, COUNT(*) as count - FROM resource_allocations - GROUP BY resource_type, status - """) - - # Get recent activity - recent_activity = await conn.fetch(""" - SELECT ch.action, ch.timestamp, cr.title, cr.requesting_product - FROM coordination_history ch - JOIN coordination_requests cr ON ch.coordination_request_id = cr.id - ORDER BY ch.timestamp DESC - LIMIT 10 - """) - - return { - "system_status": "operational", - "registered_products": product_count, - "coordination_request_stats": [dict(stat) for stat in request_stats], - "resource_allocation_stats": [dict(stat) for stat in resource_stats], - "recent_activity": [dict(activity) for activity in recent_activity], - "last_updated": datetime.utcnow().isoformat(), - } - - except Exception as e: - logger.error(f"Error getting coordination status: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to get coordination status: {str(e)}" - ) - - -@router.get("/protocols", summary="List coordination protocols") -async def list_coordination_protocols(): - """List available coordination protocols and procedures""" - try: - async with get_db_connection() as conn: - protocols = await conn.fetch(""" - SELECT protocol_name, coordination_type, scope, procedure_steps, - required_approvals, sla_requirements, is_active, version - FROM coordination_protocols - WHERE is_active = true - ORDER BY protocol_name ASC - """) - - return { - "coordination_protocols": [dict(protocol) for protocol in protocols], - "total_count": len(protocols), - } - - except Exception as e: - logger.error(f"Error listing coordination protocols: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to list coordination protocols: {str(e)}" - ) +""" +Cross-Product Coordination API Endpoints + +This module provides REST API endpoints for managing cross-product coordination +within the WCG ecosystem. Enables centralized coordination between FuzeAgent, +FuzeFront, HubHit, DeployAI, and other WCG products. +""" + +import logging +from datetime import date, datetime +from typing import Any, Dict, List, Optional + +import asyncpg +from fastapi import APIRouter, Depends, HTTPException, Path, Query +from pydantic import BaseModel, Field + +from .database import get_db_connection + +logger = logging.getLogger(__name__) +router = APIRouter(prefix="/coordination", tags=["Cross-Product Coordination"]) + + +# Pydantic Models +class ProductRegistration(BaseModel): + id: str = Field(..., description="Unique product identifier") + name: str = Field(..., description="Product display name") + version: str = Field(..., description="Current product version") + endpoints: List[str] = Field(default=[], description="API endpoints exposed") + dependencies: List[str] = Field(default=[], description="Product dependencies") + resource_requirements: Dict[str, Any] = Field( + default={}, description="Resource needs" + ) + team_contacts: List[str] = Field(default=[], description="Team contact information") + priority_level: int = Field(default=5, ge=1, le=10, description="Business priority") + metadata: Dict[str, Any] = Field(default={}, description="Additional metadata") + + +class CoordinationRequestCreate(BaseModel): + requesting_product: str = Field(..., description="Product making the request") + target_products: List[str] = Field( + ..., description="Target products for coordination" + ) + coordination_type: str = Field(..., description="Type of coordination needed") + scope: str = Field(default="product_group", description="Coordination scope") + priority: str = Field(default="medium", description="Request priority") + title: str = Field(..., description="Brief title for the request") + description: str = Field(..., description="Detailed description") + resource_requirements: Dict[str, Any] = Field( + default={}, description="Required resources" + ) + proposed_timeline: Dict[str, Any] = Field( + default={}, description="Proposed timeline" + ) + + +class CoordinationResolution(BaseModel): + resolution_plan: Dict[str, Any] = Field(..., description="Detailed resolution plan") + resolver_id: str = Field(..., description="ID of the resolver") + notes: Optional[str] = Field(None, description="Additional resolution notes") + + +class ResourceAllocationCreate(BaseModel): + product_id: str = Field(..., description="Product requesting allocation") + resource_type: str = Field(..., description="Type of resource") + resource_name: str = Field(..., description="Specific resource name") + allocation_details: Dict[str, Any] = Field( + default={}, description="Allocation specifics" + ) + valid_until: Optional[datetime] = Field(None, description="Allocation expiry") + + +# Product Registration Endpoints +@router.post( + "/products/register", summary="Register a new product in coordination system" +) +async def register_product(product: ProductRegistration): + """Register a new product in the cross-product coordination system""" + try: + async with get_db_connection() as conn: + await conn.execute( + """ + INSERT INTO product_registry ( + id, name, version, endpoints, dependencies, + resource_requirements, team_contacts, priority_level, + metadata, registered_at, updated_at + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) + ON CONFLICT (id) DO UPDATE SET + name = $2, version = $3, endpoints = $4, dependencies = $5, + resource_requirements = $6, team_contacts = $7, + priority_level = $8, metadata = $9, updated_at = $11 + """, + product.id, + product.name, + product.version, + product.endpoints, + product.dependencies, + product.resource_requirements, + product.team_contacts, + product.priority_level, + product.metadata, + datetime.utcnow(), + datetime.utcnow(), + ) + + return { + "status": "success", + "product_id": product.id, + "message": "Product registered successfully", + } + + except Exception as e: + logger.error(f"Error registering product {product.id}: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to register product: {str(e)}" + ) + + +@router.get("/products", summary="List all registered products") +async def list_products( + priority_min: int = Query(1, ge=1, le=10, description="Minimum priority level"), + active_only: bool = Query(True, description="Show only active products"), +): + """List all products registered in the coordination system""" + try: + async with get_db_connection() as conn: + products = await conn.fetch( + """ + SELECT id, name, version, priority_level, metadata, + endpoints, dependencies, registered_at, updated_at + FROM product_registry + WHERE priority_level >= $1 + ORDER BY priority_level DESC, name ASC + """, + priority_min, + ) + + return { + "products": [dict(product) for product in products], + "total_count": len(products), + } + + except Exception as e: + logger.error(f"Error listing products: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to list products: {str(e)}" + ) + + +@router.get("/products/{product_id}", summary="Get specific product details") +async def get_product(product_id: str = Path(..., description="Product ID")): + """Get detailed information about a specific product""" + try: + async with get_db_connection() as conn: + product = await conn.fetchrow( + """ + SELECT * FROM product_registry WHERE id = $1 + """, + product_id, + ) + + if not product: + raise HTTPException( + status_code=404, detail=f"Product {product_id} not found" + ) + + # Get active coordination requests involving this product + async with get_db_connection() as conn: + coordination_requests = await conn.fetch( + """ + SELECT id, title, coordination_type, priority, status, created_at + FROM coordination_requests + WHERE requesting_product = $1 + OR $1 = ANY(string_to_array(replace(replace(target_products::text, '[', ''), ']', ''), ',')) + ORDER BY created_at DESC LIMIT 10 + """, + product_id, + ) + + return { + "product": dict(product), + "active_coordination_requests": [ + dict(req) for req in coordination_requests + ], + } + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting product {product_id}: {e}") + raise HTTPException(status_code=500, detail=f"Failed to get product: {str(e)}") + + +# Coordination Request Endpoints +@router.post("/requests", summary="Create a new coordination request") +async def create_coordination_request(request: CoordinationRequestCreate): + """Create a new cross-product coordination request""" + try: + # Validate requesting product exists + async with get_db_connection() as conn: + requesting_product = await conn.fetchrow( + """ + SELECT id FROM product_registry WHERE id = $1 + """, + request.requesting_product, + ) + + if not requesting_product: + raise HTTPException( + status_code=400, + detail=f"Requesting product {request.requesting_product} not found", + ) + + # Create coordination request + async with get_db_connection() as conn: + request_id = await conn.fetchval( + """ + INSERT INTO coordination_requests ( + requesting_product, target_products, coordination_type, + scope, priority, title, description, resource_requirements, + proposed_timeline, stakeholders, created_at, updated_at + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) + RETURNING id + """, + request.requesting_product, + request.target_products, + request.coordination_type, + request.scope, + request.priority, + request.title, + request.description, + request.resource_requirements, + request.proposed_timeline, + [], # stakeholders - could be auto-populated + datetime.utcnow(), + datetime.utcnow(), + ) + + # Log coordination history + async with get_db_connection() as conn: + await conn.execute( + """ + INSERT INTO coordination_history ( + coordination_request_id, action, actor_type, details + ) VALUES ($1, $2, $3, $4) + """, + request_id, + "created", + "system", + {"created_via": "api"}, + ) + + return { + "status": "success", + "request_id": str(request_id), + "message": "Coordination request created successfully", + } + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error creating coordination request: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to create coordination request: {str(e)}" + ) + + +@router.get("/requests", summary="List coordination requests") +async def list_coordination_requests( + status: Optional[str] = Query(None, description="Filter by status"), + priority: Optional[str] = Query(None, description="Filter by priority"), + product_id: Optional[str] = Query(None, description="Filter by product"), + limit: int = Query(50, ge=1, le=200, description="Max number of results"), +): + """List coordination requests with optional filters""" + try: + where_conditions = [] + params = [] + param_count = 0 + + if status: + param_count += 1 + where_conditions.append(f"status = ${param_count}") + params.append(status) + + if priority: + param_count += 1 + where_conditions.append(f"priority = ${param_count}") + params.append(priority) + + if product_id: + param_count += 1 + where_conditions.append( + f"(requesting_product = ${param_count} OR ${param_count} = ANY(string_to_array(replace(replace(target_products::text, '[', ''), ']', ''), ',')))" + ) + params.append(product_id) + + where_clause = ( + " WHERE " + " AND ".join(where_conditions) if where_conditions else "" + ) + param_count += 1 + params.append(limit) + + query = f""" + SELECT id, requesting_product, target_products, coordination_type, + scope, priority, status, title, description, created_at, updated_at + FROM coordination_requests + {where_clause} + ORDER BY + CASE priority + WHEN 'critical' THEN 1 + WHEN 'high' THEN 2 + WHEN 'medium' THEN 3 + WHEN 'low' THEN 4 + END, + created_at DESC + LIMIT ${param_count} + """ # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params + + async with get_db_connection() as conn: + requests = await conn.fetch(query, *params) + + return { + "coordination_requests": [dict(req) for req in requests], + "total_count": len(requests), + } + + except Exception as e: + logger.error(f"Error listing coordination requests: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to list coordination requests: {str(e)}" + ) + + +@router.get("/requests/{request_id}", summary="Get coordination request details") +async def get_coordination_request( + request_id: str = Path(..., description="Coordination request ID") +): + """Get detailed information about a specific coordination request""" + try: + async with get_db_connection() as conn: + request = await conn.fetchrow( + """ + SELECT * FROM coordination_requests WHERE id = $1 + """, + request_id, + ) + + if not request: + raise HTTPException( + status_code=404, detail=f"Coordination request {request_id} not found" + ) + + # Get coordination history + async with get_db_connection() as conn: + history = await conn.fetch( + """ + SELECT action, actor_id, actor_type, details, timestamp + FROM coordination_history + WHERE coordination_request_id = $1 + ORDER BY timestamp ASC + """, + request_id, + ) + + return { + "coordination_request": dict(request), + "history": [dict(h) for h in history], + } + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting coordination request {request_id}: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to get coordination request: {str(e)}" + ) + + +@router.put("/requests/{request_id}/resolve", summary="Resolve coordination request") +async def resolve_coordination_request( + request_id: str = Path(..., description="Coordination request ID"), + resolution: CoordinationResolution = ..., +): + """Resolve a coordination request with a specific plan""" + try: + async with get_db_connection() as conn: + # Check if request exists and is pending + existing_request = await conn.fetchrow( + """ + SELECT id, status FROM coordination_requests WHERE id = $1 + """, + request_id, + ) + + if not existing_request: + raise HTTPException( + status_code=404, detail=f"Coordination request {request_id} not found" + ) + + if existing_request["status"] not in ["pending", "in_progress"]: + raise HTTPException( + status_code=400, + detail=f"Request is already {existing_request['status']}", + ) + + # Update request status and resolution + async with get_db_connection() as conn: + await conn.execute( + """ + UPDATE coordination_requests SET + status = 'resolved', + resolution_plan = $2, + resolved_at = $3, + updated_at = $4 + WHERE id = $1 + """, + request_id, + resolution.resolution_plan, + datetime.utcnow(), + datetime.utcnow(), + ) + + # Log resolution in history + async with get_db_connection() as conn: + await conn.execute( + """ + INSERT INTO coordination_history ( + coordination_request_id, action, actor_id, actor_type, details + ) VALUES ($1, $2, $3, $4, $5) + """, + request_id, + "resolved", + resolution.resolver_id, + "agent", + { + "resolution_plan": resolution.resolution_plan, + "notes": resolution.notes, + }, + ) + + return { + "status": "success", + "request_id": request_id, + "message": "Coordination request resolved successfully", + } + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error resolving coordination request {request_id}: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to resolve coordination request: {str(e)}" + ) + + +# Resource Management Endpoints +@router.post("/resources/allocate", summary="Allocate resources to a product") +async def allocate_resource(allocation: ResourceAllocationCreate): + """Allocate a resource to a specific product""" + try: + async with get_db_connection() as conn: + allocation_id = await conn.fetchval( + """ + INSERT INTO resource_allocations ( + product_id, resource_type, resource_name, allocation_details, + valid_until, status, created_at, updated_at + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + RETURNING id + """, + allocation.product_id, + allocation.resource_type, + allocation.resource_name, + allocation.allocation_details, + allocation.valid_until, + "active", + datetime.utcnow(), + datetime.utcnow(), + ) + + return { + "status": "success", + "allocation_id": str(allocation_id), + "message": "Resource allocated successfully", + } + + except asyncpg.UniqueViolationError: + raise HTTPException( + status_code=409, + detail=f"Resource {allocation.resource_name} of type {allocation.resource_type} is already allocated", + ) + except Exception as e: + logger.error(f"Error allocating resource: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to allocate resource: {str(e)}" + ) + + +@router.get("/resources", summary="List resource allocations") +async def list_resource_allocations( + product_id: Optional[str] = Query(None, description="Filter by product"), + resource_type: Optional[str] = Query(None, description="Filter by resource type"), + status: Optional[str] = Query(None, description="Filter by status"), +): + """List current resource allocations""" + try: + where_conditions = [] + params = [] + param_count = 0 + + if product_id: + param_count += 1 + where_conditions.append(f"product_id = ${param_count}") + params.append(product_id) + + if resource_type: + param_count += 1 + where_conditions.append(f"resource_type = ${param_count}") + params.append(resource_type) + + if status: + param_count += 1 + where_conditions.append(f"status = ${param_count}") + params.append(status) + + where_clause = ( + " WHERE " + " AND ".join(where_conditions) if where_conditions else "" + ) + + query = f""" + SELECT ra.*, pr.name as product_name + FROM resource_allocations ra + LEFT JOIN product_registry pr ON ra.product_id = pr.id + {where_clause} + ORDER BY ra.created_at DESC + """ # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params + + async with get_db_connection() as conn: + allocations = await conn.fetch(query, *params) + + return { + "resource_allocations": [dict(alloc) for alloc in allocations], + "total_count": len(allocations), + } + + except Exception as e: + logger.error(f"Error listing resource allocations: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to list resource allocations: {str(e)}" + ) + + +# Dashboard and Status Endpoints +@router.get("/status", summary="Get overall coordination system status") +async def get_coordination_status(): + """Get overall status of the cross-product coordination system""" + try: + async with get_db_connection() as conn: + # Get request counts by status + request_stats = await conn.fetch(""" + SELECT status, priority, COUNT(*) as count + FROM coordination_requests + GROUP BY status, priority + ORDER BY status, priority + """) + + # Get product count + product_count = await conn.fetchval(""" + SELECT COUNT(*) FROM product_registry + """) + + # Get resource allocation stats + resource_stats = await conn.fetch(""" + SELECT resource_type, status, COUNT(*) as count + FROM resource_allocations + GROUP BY resource_type, status + """) + + # Get recent activity + recent_activity = await conn.fetch(""" + SELECT ch.action, ch.timestamp, cr.title, cr.requesting_product + FROM coordination_history ch + JOIN coordination_requests cr ON ch.coordination_request_id = cr.id + ORDER BY ch.timestamp DESC + LIMIT 10 + """) + + return { + "system_status": "operational", + "registered_products": product_count, + "coordination_request_stats": [dict(stat) for stat in request_stats], + "resource_allocation_stats": [dict(stat) for stat in resource_stats], + "recent_activity": [dict(activity) for activity in recent_activity], + "last_updated": datetime.utcnow().isoformat(), + } + + except Exception as e: + logger.error(f"Error getting coordination status: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to get coordination status: {str(e)}" + ) + + +@router.get("/protocols", summary="List coordination protocols") +async def list_coordination_protocols(): + """List available coordination protocols and procedures""" + try: + async with get_db_connection() as conn: + protocols = await conn.fetch(""" + SELECT protocol_name, coordination_type, scope, procedure_steps, + required_approvals, sla_requirements, is_active, version + FROM coordination_protocols + WHERE is_active = true + ORDER BY protocol_name ASC + """) + + return { + "coordination_protocols": [dict(protocol) for protocol in protocols], + "total_count": len(protocols), + } + + except Exception as e: + logger.error(f"Error listing coordination protocols: {e}") + raise HTTPException( + status_code=500, detail=f"Failed to list coordination protocols: {str(e)}" + ) diff --git a/services/orchestrator/goal_conversation_service.py b/services/orchestrator/goal_conversation_service.py index ce16acf..bf2be4f 100644 --- a/services/orchestrator/goal_conversation_service.py +++ b/services/orchestrator/goal_conversation_service.py @@ -1,1040 +1,1040 @@ -""" -Goal Conversation Management Service for FuzeAgent - -This service manages AI-powered conversations about organizational goals, -enabling collaborative planning, progress reviews, problem-solving, and -strategic adjustments through intelligent dialogue. -""" - -import asyncio -import json -import logging -import uuid -from dataclasses import dataclass -from datetime import datetime, timedelta -from enum import Enum -from typing import Any, Dict, List, Optional, Tuple - -import asyncpg - -logger = logging.getLogger(__name__) - - -class ConversationType(str, Enum): - PLANNING = "planning" - REVIEW = "review" - ADJUSTMENT = "adjustment" - PROBLEM_SOLVING = "problem_solving" - BRAINSTORMING = "brainstorming" - RETROSPECTIVE = "retrospective" - - -class ConversationStatus(str, Enum): - ACTIVE = "active" - ARCHIVED = "archived" - COMPLETED = "completed" - - -class MessageType(str, Enum): - SYSTEM = "system" - AGENT = "agent" - HUMAN = "human" - AI_ANALYSIS = "ai_analysis" - ACTION_ITEM = "action_item" - - -@dataclass -class ConversationMessage: - """Represents a message in a goal conversation""" - - id: str - message_type: MessageType - sender_id: Optional[str] - sender_name: Optional[str] - content: str - metadata: Dict[str, Any] - timestamp: datetime - references: List[str] # Referenced message IDs - reactions: List[Dict[str, Any]] # Message reactions/acknowledgments - - -@dataclass -class ConversationInsight: - """Represents an AI-generated insight from conversation analysis""" - - id: str - insight_type: str # pattern, risk, opportunity, recommendation - title: str - description: str - confidence_score: float - supporting_messages: List[str] - suggested_actions: List[Dict[str, Any]] - generated_at: datetime - - -@dataclass -class ActionItem: - """Represents an action item derived from conversation""" - - id: str - title: str - description: str - assigned_to: Optional[str] - due_date: Optional[datetime] - status: str # pending, in_progress, completed, cancelled - priority: int - source_messages: List[str] - created_at: datetime - completed_at: Optional[datetime] - - -class GoalConversationService: - """ - Manages AI-powered conversations for organizational goal planning, - tracking, and optimization with intelligent insights and action generation. - """ - - def __init__(self, database_url: str): - self.database_url = database_url - self.pool: Optional[asyncpg.Pool] = None - - # Configuration - self.max_conversation_messages = 1000 - self.insight_confidence_threshold = 0.6 - self.auto_action_item_threshold = 0.8 - - # AI conversation templates and prompts - self.conversation_starters = self._initialize_conversation_starters() - self.analysis_prompts = self._initialize_analysis_prompts() - - # Statistics - self.conversations_created = 0 - self.messages_processed = 0 - self.insights_generated = 0 - self.action_items_created = 0 - - async def initialize(self): - """Initialize the goal conversation service""" - logger.info("Initializing GoalConversationService") - - try: - self.pool = await asyncpg.create_pool( - self.database_url, min_size=1, max_size=5, command_timeout=60 - ) - - logger.info("GoalConversationService initialized successfully") - - except Exception as e: - logger.error(f"Failed to initialize GoalConversationService: {e}") - raise - - async def close(self): - """Close database connections""" - if self.pool: - await self.pool.close() - logger.info("GoalConversationService closed") - - async def create_goal_conversation( - self, - goal_id: str, - conversation_type: ConversationType, - conversation_title: str, - initial_context: Optional[Dict[str, Any]] = None, - participants: Optional[List[Dict[str, Any]]] = None, - created_by: Optional[str] = None, - ) -> str: - """Create a new conversation for a goal""" - - conversation_id = str(uuid.uuid4()) - - if initial_context is None: - initial_context = {} - if participants is None: - participants = [] - - try: - async with self.pool.acquire() as conn: - # Get goal context - goal = await conn.fetchrow( - """ - SELECT title, description, goal_type, target_deadline, - progress_percentage, current_value, target_value - FROM organization_goals WHERE id = $1 - """, - goal_id, - ) - - if not goal: - raise ValueError(f"Goal {goal_id} not found") - - # Enhanced context with goal information - enhanced_context = { - **initial_context, - "goal_title": goal["title"], - "goal_type": goal["goal_type"], - "goal_progress": float(goal["progress_percentage"]), - "days_to_deadline": ( - goal["target_deadline"] - datetime.now().date() - ).days, - "conversation_created_at": datetime.now().isoformat(), - } - - # Create conversation - await conn.execute( - """ - INSERT INTO goal_conversations ( - id, goal_id, conversation_type, conversation_title, - conversation_context, participants, status, created_by - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) - """, - conversation_id, - goal_id, - conversation_type.value, - conversation_title, - json.dumps(enhanced_context), - json.dumps(participants), - ConversationStatus.ACTIVE.value, - created_by, - ) - - # Add initial system message with conversation starter - starter_message = self._generate_conversation_starter( - conversation_type, goal, enhanced_context - ) - - await self._add_message( - conversation_id, - MessageType.SYSTEM, - None, - "System", - starter_message, - {"conversation_starter": True}, - ) - - self.conversations_created += 1 - logger.info(f"Created conversation {conversation_id} for goal {goal_id}") - - return conversation_id - - except Exception as e: - logger.error(f"Error creating conversation: {e}") - raise - - async def add_message_to_conversation( - self, - conversation_id: str, - message_type: MessageType, - sender_id: Optional[str], - sender_name: str, - content: str, - metadata: Optional[Dict[str, Any]] = None, - references: Optional[List[str]] = None, - ) -> str: - """Add a message to a conversation""" - - if metadata is None: - metadata = {} - if references is None: - references = [] - - try: - # Add the message - message_id = await self._add_message( - conversation_id, - message_type, - sender_id, - sender_name, - content, - metadata, - references, - ) - - # Trigger conversation analysis for insights - await self._analyze_conversation_for_insights(conversation_id) - - # Update conversation activity timestamp - async with self.pool.acquire() as conn: - await conn.execute( - """ - UPDATE goal_conversations - SET last_activity_at = NOW(), updated_at = NOW() - WHERE id = $1 - """, - conversation_id, - ) - - self.messages_processed += 1 - - return message_id - - except Exception as e: - logger.error(f"Error adding message to conversation {conversation_id}: {e}") - raise - - async def get_conversation(self, conversation_id: str) -> Optional[Dict[str, Any]]: - """Get full conversation with messages, insights, and action items""" - - try: - async with self.pool.acquire() as conn: - # Get conversation details - conversation = await conn.fetchrow( - """ - SELECT gc.*, og.title as goal_title, og.goal_type - FROM goal_conversations gc - JOIN organization_goals og ON gc.goal_id = og.id - WHERE gc.id = $1 - """, - conversation_id, - ) - - if not conversation: - return None - - # Get messages - messages = ( - json.loads(conversation["messages"]) - if conversation["messages"] - else [] - ) - - # Get insights - insights = ( - json.loads(conversation["insights_generated"]) - if conversation["insights_generated"] - else [] - ) - - # Get action items - action_items = ( - json.loads(conversation["action_items"]) - if conversation["action_items"] - else [] - ) - - return { - "id": str(conversation["id"]), - "goal_id": str(conversation["goal_id"]), - "goal_title": conversation["goal_title"], - "conversation_type": conversation["conversation_type"], - "conversation_title": conversation["conversation_title"], - "conversation_summary": conversation["conversation_summary"], - "conversation_context": ( - json.loads(conversation["conversation_context"]) - if conversation["conversation_context"] - else {} - ), - "participants": ( - json.loads(conversation["participants"]) - if conversation["participants"] - else [] - ), - "messages": messages, - "insights_generated": insights, - "action_items": action_items, - "status": conversation["status"], - "last_activity_at": ( - conversation["last_activity_at"].isoformat() - if conversation["last_activity_at"] - else None - ), - "created_at": conversation["created_at"].isoformat(), - "updated_at": conversation["updated_at"].isoformat(), - "message_count": len(messages), - "insight_count": len(insights), - "action_item_count": len(action_items), - } - - except Exception as e: - logger.error(f"Error getting conversation {conversation_id}: {e}") - return None - - async def generate_planning_milestones( - self, conversation_id: str, planning_context: Optional[Dict[str, Any]] = None - ) -> List[Dict[str, Any]]: - """Generate milestone suggestions based on conversation analysis""" - - try: - async with self.pool.acquire() as conn: - # Get conversation and goal context - conversation = await conn.fetchrow( - """ - SELECT gc.*, og.title, og.description, og.goal_type, - og.target_deadline, og.target_value, og.target_unit - FROM goal_conversations gc - JOIN organization_goals og ON gc.goal_id = og.id - WHERE gc.id = $1 - """, - conversation_id, - ) - - if not conversation: - raise ValueError(f"Conversation {conversation_id} not found") - - # Analyze conversation content for milestone ideas - messages = ( - json.loads(conversation["messages"]) - if conversation["messages"] - else [] - ) - milestone_suggestions = self._extract_milestone_ideas_from_conversation( - messages, conversation, planning_context - ) - - # Generate AI-powered milestone recommendations - ai_milestones = await self._generate_ai_milestone_recommendations( - conversation, milestone_suggestions, planning_context - ) - - # Add milestones as insights to the conversation - milestone_insight = { - "id": str(uuid.uuid4()), - "insight_type": "milestone_recommendations", - "title": "AI-Generated Milestone Recommendations", - "description": f"Based on conversation analysis, here are {len(ai_milestones)} recommended milestones", - "confidence_score": 0.85, - "supporting_messages": [ - msg["id"] for msg in messages[-5:] if "id" in msg - ], # Last 5 messages - "suggested_actions": [ - { - "action": "create_milestones", - "description": "Create these milestones for the goal", - "milestones": ai_milestones, - } - ], - "generated_at": datetime.now().isoformat(), - } - - # Update conversation with milestone insight - await self._add_insight_to_conversation( - conversation_id, milestone_insight - ) - - return ai_milestones - - except Exception as e: - logger.error(f"Error generating planning milestones: {e}") - return [] - - async def conduct_progress_review( - self, conversation_id: str, review_period_days: int = 30 - ) -> Dict[str, Any]: - """Conduct AI-powered progress review for a goal conversation""" - - try: - async with self.pool.acquire() as conn: - # Get conversation and goal data - conversation = await conn.fetchrow( - """ - SELECT gc.*, og.* - FROM goal_conversations gc - JOIN organization_goals og ON gc.goal_id = og.id - WHERE gc.id = $1 - """, - conversation_id, - ) - - if not conversation: - raise ValueError(f"Conversation {conversation_id} not found") - - # Get recent progress data - progress_data = await conn.fetch( - """ - SELECT * FROM goal_progress_tracking - WHERE goal_id = $1 - AND recorded_at >= NOW() - INTERVAL '%s days' - ORDER BY recorded_at DESC - """, - str(conversation["goal_id"]), - review_period_days, - ) - - # Get milestones and tasks status - milestone_status = await conn.fetch( - """ - SELECT status, COUNT(*) as count - FROM goal_milestones - WHERE goal_id = $1 - GROUP BY status - """, - str(conversation["goal_id"]), - ) - - task_status = await conn.fetch( - """ - SELECT status, COUNT(*) as count - FROM goal_tasks - WHERE goal_id = $1 - GROUP BY status - """, - str(conversation["goal_id"]), - ) - - # Generate review analysis - review_analysis = { - "review_period_days": review_period_days, - "goal_progress": { - "current_progress": float(conversation["progress_percentage"]), - "target_value": ( - float(conversation["target_value"]) - if conversation["target_value"] - else None - ), - "current_value": ( - float(conversation["current_value"]) - if conversation["current_value"] - else None - ), - "completion_confidence": float( - conversation["completion_confidence"] - ), - }, - "milestone_summary": { - row["status"]: row["count"] for row in milestone_status - }, - "task_summary": { - row["status"]: row["count"] for row in task_status - }, - "progress_trend": self._calculate_progress_trend( - [dict(p) for p in progress_data] - ), - "risk_assessment": self._assess_goal_risks( - conversation, progress_data - ), - "recommendations": self._generate_progress_recommendations( - conversation, progress_data - ), - } - - # Add review as a structured message - review_message = self._format_progress_review_message(review_analysis) - await self._add_message( - conversation_id, - MessageType.AI_ANALYSIS, - None, - "Progress Analyzer", - review_message, - {"review_analysis": review_analysis}, - ) - - return review_analysis - - except Exception as e: - logger.error(f"Error conducting progress review: {e}") - return {"error": str(e)} - - async def extract_action_items_from_conversation( - self, conversation_id: str, auto_assign: bool = True - ) -> List[Dict[str, Any]]: - """Extract and create action items from conversation analysis""" - - try: - async with self.pool.acquire() as conn: - conversation = await conn.fetchrow( - """ - SELECT * FROM goal_conversations WHERE id = $1 - """, - conversation_id, - ) - - if not conversation: - raise ValueError(f"Conversation {conversation_id} not found") - - messages = ( - json.loads(conversation["messages"]) - if conversation["messages"] - else [] - ) - - # Analyze messages for actionable items - potential_actions = self._identify_action_items_in_messages(messages) - - # Convert to action item format - action_items = [] - for action in potential_actions: - if action["confidence"] >= self.auto_action_item_threshold: - action_item = { - "id": str(uuid.uuid4()), - "title": action["title"], - "description": action["description"], - "assigned_to": ( - action.get("assigned_to") if auto_assign else None - ), - "due_date": action.get("due_date"), - "status": "pending", - "priority": action.get("priority", 5), - "source_messages": action["source_messages"], - "created_at": datetime.now().isoformat(), - "confidence_score": action["confidence"], - } - action_items.append(action_item) - - # Update conversation with action items - if action_items: - existing_actions = ( - json.loads(conversation["action_items"]) - if conversation["action_items"] - else [] - ) - all_actions = existing_actions + action_items - - await conn.execute( - """ - UPDATE goal_conversations - SET action_items = $2, updated_at = NOW() - WHERE id = $1 - """, - conversation_id, - json.dumps(all_actions), - ) - - self.action_items_created += len(action_items) - - return action_items - - except Exception as e: - logger.error(f"Error extracting action items: {e}") - return [] - - async def get_goal_conversations( - self, - goal_id: str, - conversation_type: Optional[ConversationType] = None, - status: Optional[ConversationStatus] = None, - limit: int = 10, - ) -> List[Dict[str, Any]]: - """Get conversations for a goal with optional filtering""" - - try: - async with self.pool.acquire() as conn: - where_conditions = ["goal_id = $1"] - params = [goal_id] - param_idx = 2 - - if conversation_type: - where_conditions.append(f"conversation_type = ${param_idx}") - params.append(conversation_type.value) - param_idx += 1 - - if status: - where_conditions.append(f"status = ${param_idx}") - params.append(status.value) - param_idx += 1 - - where_clause = " AND ".join(where_conditions) - - conversations = await conn.fetch( - """ - SELECT id, conversation_type, conversation_title, conversation_summary, - status, last_activity_at, created_at, - COALESCE(array_length(string_to_array(messages::text, '}}'), 1), 0) as message_count, - COALESCE(array_length(string_to_array(action_items::text, '}}'), 1), 0) as action_count - FROM goal_conversations - WHERE {where_clause} - ORDER BY last_activity_at DESC, created_at DESC - LIMIT ${param_idx} - """.format( # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params - where_clause=where_clause, param_idx=param_idx - ), - *params, - limit, - ) - - return [dict(conv) for conv in conversations] - - except Exception as e: - logger.error(f"Error getting conversations for goal {goal_id}: {e}") - return [] - - # Helper methods for conversation management - - async def _add_message( - self, - conversation_id: str, - message_type: MessageType, - sender_id: Optional[str], - sender_name: str, - content: str, - metadata: Dict[str, Any], - references: Optional[List[str]] = None, - ) -> str: - """Add a message to conversation""" - - message_id = str(uuid.uuid4()) - message = { - "id": message_id, - "message_type": message_type.value, - "sender_id": sender_id, - "sender_name": sender_name, - "content": content, - "metadata": metadata, - "timestamp": datetime.now().isoformat(), - "references": references or [], - "reactions": [], - } - - async with self.pool.acquire() as conn: - # Get current messages - current_messages = await conn.fetchval( - """ - SELECT messages FROM goal_conversations WHERE id = $1 - """, - conversation_id, - ) - - messages = json.loads(current_messages) if current_messages else [] - messages.append(message) - - # Limit message history - if len(messages) > self.max_conversation_messages: - messages = messages[-self.max_conversation_messages :] - - # Update conversation - await conn.execute( - """ - UPDATE goal_conversations - SET messages = $2, updated_at = NOW() - WHERE id = $1 - """, - conversation_id, - json.dumps(messages), - ) - - return message_id - - async def _add_insight_to_conversation( - self, conversation_id: str, insight: Dict[str, Any] - ): - """Add an insight to conversation""" - - async with self.pool.acquire() as conn: - # Get current insights - current_insights = await conn.fetchval( - """ - SELECT insights_generated FROM goal_conversations WHERE id = $1 - """, - conversation_id, - ) - - insights = json.loads(current_insights) if current_insights else [] - insights.append(insight) - - # Update conversation - await conn.execute( - """ - UPDATE goal_conversations - SET insights_generated = $2, updated_at = NOW() - WHERE id = $1 - """, - conversation_id, - json.dumps(insights), - ) - - self.insights_generated += 1 - - def _generate_conversation_starter( - self, - conversation_type: ConversationType, - goal: Dict[str, Any], - context: Dict[str, Any], - ) -> str: - """Generate an appropriate conversation starter""" - - starters = self.conversation_starters.get(conversation_type, {}) - - if conversation_type == ConversationType.PLANNING: - return starters["default"].format( - goal_title=goal["title"], - goal_type=goal["goal_type"], - deadline=goal["target_deadline"].strftime("%B %d, %Y"), - target_value=( - goal["target_value"] - if goal["target_value"] - else "defined objectives" - ), - ) - elif conversation_type == ConversationType.REVIEW: - return starters["default"].format( - goal_title=goal["title"], - current_progress=f"{goal['progress_percentage']:.1f}%", - ) - else: - return starters.get("default", f"Let's discuss the {goal['title']} goal.") - - def _initialize_conversation_starters(self) -> Dict[str, Dict[str, str]]: - """Initialize conversation starter templates""" - - return { - ConversationType.PLANNING: { - "default": """Welcome to the strategic planning session for "{goal_title}"! - -🎯 **Goal**: {goal_title} -📊 **Type**: {goal_type} -📅 **Deadline**: {deadline} -🎌 **Target**: {target_value} - -Let's break this goal down into actionable milestones and tasks. Here are some questions to get us started: - -1. **What are the major milestones we need to achieve?** -2. **What dependencies and blockers should we consider?** -3. **Which teams and resources will be involved?** -4. **How should we measure progress along the way?** - -What aspect would you like to focus on first?""" - }, - ConversationType.REVIEW: { - "default": """Time for a progress review of "{goal_title}"! - -📈 **Current Progress**: {current_progress} - -Let's evaluate our progress, identify what's working well, and address any challenges. - -Key areas to discuss: -- Recent achievements and wins -- Current blockers or risks -- Resource allocation and team performance -- Timeline adjustments if needed -- Next steps and priorities - -What would you like to review first?""" - }, - ConversationType.PROBLEM_SOLVING: { - "default": """Problem-solving session for "{goal_title}". - -Let's identify the specific challenges we're facing and work together to find solutions. Please share: -- What specific problems or blockers have emerged? -- What have we tried so far? -- What constraints or requirements should we consider? - -What's the main challenge you'd like to tackle?""" - }, - } - - def _initialize_analysis_prompts(self) -> Dict[str, str]: - """Initialize AI analysis prompts for conversation processing""" - - return { - "extract_milestones": """Analyze this goal conversation and extract potential milestones mentioned or implied. Look for: -- Time-based deliverables or checkpoints -- Measurable objectives or targets -- Dependencies between activities -- Key decision points or reviews - -Return milestones with titles, descriptions, and target dates.""", - "identify_risks": """Analyze this conversation for potential risks, blockers, or concerns mentioned. Look for: -- Resource constraints or availability issues -- Technical challenges or unknowns -- Timeline concerns or dependencies -- Team capacity or skill gaps -- External dependencies or market factors - -Assess the probability and impact of each risk.""", - "extract_actions": """Extract specific action items from this conversation. Look for: -- Tasks or activities that someone needs to do -- Decisions that need to be made -- Information that needs to be gathered -- People who need to be contacted -- Deadlines or time-sensitive items - -Include who should be responsible and when it should be done.""", - } - - # Placeholder implementations for AI analysis methods - - async def _analyze_conversation_for_insights(self, conversation_id: str): - """Analyze conversation for insights (placeholder for AI integration)""" - # This would integrate with an AI service for conversation analysis - pass - - def _extract_milestone_ideas_from_conversation( - self, - messages: List[Dict[str, Any]], - conversation: Dict[str, Any], - planning_context: Optional[Dict[str, Any]], - ) -> List[Dict[str, Any]]: - """Extract milestone ideas from conversation messages""" - # Simplified implementation - would use NLP/AI for real extraction - milestone_keywords = [ - "milestone", - "phase", - "deliverable", - "target", - "deadline", - "complete", - ] - - milestones = [] - for message in messages: - content = message.get("content", "").lower() - if any(keyword in content for keyword in milestone_keywords): - # Extract potential milestone (simplified) - milestones.append( - { - "title": f"Milestone from conversation", - "description": message.get("content", "")[:200], - "source_message": message.get("id"), - "confidence": 0.7, - } - ) - - return milestones[:5] # Return top 5 candidates - - async def _generate_ai_milestone_recommendations( - self, - conversation: Dict[str, Any], - milestone_suggestions: List[Dict[str, Any]], - planning_context: Optional[Dict[str, Any]], - ) -> List[Dict[str, Any]]: - """Generate AI-powered milestone recommendations""" - # Placeholder implementation - would use AI for intelligent milestone generation - - goal_type = conversation["goal_type"] - timeline_days = (conversation["target_deadline"] - datetime.now().date()).days - - # Generate sample milestones based on goal type - if goal_type == "business" and "revenue" in conversation["title"].lower(): - return [ - { - "title": "Foundation Setup", - "description": "Establish initial infrastructure, team structure, and processes", - "target_date": (datetime.now() + timedelta(days=timeline_days // 4)) - .date() - .isoformat(), - "milestone_type": "checkpoint", - }, - { - "title": "Growth Phase Launch", - "description": "Execute marketing campaigns and sales initiatives", - "target_date": (datetime.now() + timedelta(days=timeline_days // 2)) - .date() - .isoformat(), - "milestone_type": "deliverable", - }, - { - "title": "Scale and Optimize", - "description": "Optimize processes and scale operations for target achievement", - "target_date": ( - datetime.now() + timedelta(days=timeline_days * 3 // 4) - ) - .date() - .isoformat(), - "milestone_type": "metric", - }, - ] - - return [] - - # Additional helper methods for conversation analysis - def _calculate_progress_trend(self, progress_data: List[Dict[str, Any]]) -> str: - """Calculate progress trend from historical data""" - if len(progress_data) < 2: - return "insufficient_data" - - # Simple trend calculation - recent_progress = progress_data[0]["progress_percentage"] - older_progress = progress_data[-1]["progress_percentage"] - - if recent_progress > older_progress * 1.1: - return "accelerating" - elif recent_progress < older_progress * 0.9: - return "declining" - else: - return "steady" - - def _assess_goal_risks( - self, goal: Dict[str, Any], progress_data: List[Dict[str, Any]] - ) -> List[Dict[str, Any]]: - """Assess risks based on goal and progress data""" - risks = [] - - # Timeline risk - days_remaining = (goal["target_deadline"] - datetime.now().date()).days - progress = float(goal["progress_percentage"]) - - if days_remaining < 30 and progress < 70: - risks.append( - { - "type": "timeline_risk", - "severity": "high", - "description": "Goal progress is behind schedule with limited time remaining", - } - ) - - return risks - - def _generate_progress_recommendations( - self, goal: Dict[str, Any], progress_data: List[Dict[str, Any]] - ) -> List[str]: - """Generate recommendations based on progress analysis""" - recommendations = [] - - progress = float(goal["progress_percentage"]) - - if progress < 25: - recommendations.append( - "Consider breaking down remaining work into smaller, more manageable tasks" - ) - - if progress > 75: - recommendations.append( - "Focus on final quality checks and prepare for goal completion" - ) - - return recommendations - - def _identify_action_items_in_messages( - self, messages: List[Dict[str, Any]] - ) -> List[Dict[str, Any]]: - """Identify potential action items from messages""" - # Simplified implementation - would use NLP for real extraction - action_keywords = [ - "need to", - "should", - "must", - "will", - "action", - "task", - "todo", - ] - - actions = [] - for message in messages: - content = message.get("content", "").lower() - if any(keyword in content for keyword in action_keywords): - actions.append( - { - "title": f"Action item from conversation", - "description": message.get("content", "")[:200], - "source_messages": [message.get("id")], - "confidence": 0.8, - "priority": 5, - } - ) - - return actions[:10] # Return top 10 candidates - - def _format_progress_review_message(self, review_analysis: Dict[str, Any]) -> str: - """Format progress review analysis as a readable message""" - - progress = review_analysis["goal_progress"]["current_progress"] - trend = review_analysis["progress_trend"] - - message = f"""## Progress Review Summary - -**Current Progress**: {progress:.1f}% -**Trend**: {trend.replace('_', ' ').title()} - -### Key Metrics -""" - - if review_analysis["milestone_summary"]: - message += "\n**Milestones:**\n" - for status, count in review_analysis["milestone_summary"].items(): - message += f"- {status.replace('_', ' ').title()}: {count}\n" - - if review_analysis["recommendations"]: - message += "\n### Recommendations\n" - for i, rec in enumerate(review_analysis["recommendations"], 1): - message += f"{i}. {rec}\n" - - return message +""" +Goal Conversation Management Service for FuzeAgent + +This service manages AI-powered conversations about organizational goals, +enabling collaborative planning, progress reviews, problem-solving, and +strategic adjustments through intelligent dialogue. +""" + +import asyncio +import json +import logging +import uuid +from dataclasses import dataclass +from datetime import datetime, timedelta +from enum import Enum +from typing import Any, Dict, List, Optional, Tuple + +import asyncpg + +logger = logging.getLogger(__name__) + + +class ConversationType(str, Enum): + PLANNING = "planning" + REVIEW = "review" + ADJUSTMENT = "adjustment" + PROBLEM_SOLVING = "problem_solving" + BRAINSTORMING = "brainstorming" + RETROSPECTIVE = "retrospective" + + +class ConversationStatus(str, Enum): + ACTIVE = "active" + ARCHIVED = "archived" + COMPLETED = "completed" + + +class MessageType(str, Enum): + SYSTEM = "system" + AGENT = "agent" + HUMAN = "human" + AI_ANALYSIS = "ai_analysis" + ACTION_ITEM = "action_item" + + +@dataclass +class ConversationMessage: + """Represents a message in a goal conversation""" + + id: str + message_type: MessageType + sender_id: Optional[str] + sender_name: Optional[str] + content: str + metadata: Dict[str, Any] + timestamp: datetime + references: List[str] # Referenced message IDs + reactions: List[Dict[str, Any]] # Message reactions/acknowledgments + + +@dataclass +class ConversationInsight: + """Represents an AI-generated insight from conversation analysis""" + + id: str + insight_type: str # pattern, risk, opportunity, recommendation + title: str + description: str + confidence_score: float + supporting_messages: List[str] + suggested_actions: List[Dict[str, Any]] + generated_at: datetime + + +@dataclass +class ActionItem: + """Represents an action item derived from conversation""" + + id: str + title: str + description: str + assigned_to: Optional[str] + due_date: Optional[datetime] + status: str # pending, in_progress, completed, cancelled + priority: int + source_messages: List[str] + created_at: datetime + completed_at: Optional[datetime] + + +class GoalConversationService: + """ + Manages AI-powered conversations for organizational goal planning, + tracking, and optimization with intelligent insights and action generation. + """ + + def __init__(self, database_url: str): + self.database_url = database_url + self.pool: Optional[asyncpg.Pool] = None + + # Configuration + self.max_conversation_messages = 1000 + self.insight_confidence_threshold = 0.6 + self.auto_action_item_threshold = 0.8 + + # AI conversation templates and prompts + self.conversation_starters = self._initialize_conversation_starters() + self.analysis_prompts = self._initialize_analysis_prompts() + + # Statistics + self.conversations_created = 0 + self.messages_processed = 0 + self.insights_generated = 0 + self.action_items_created = 0 + + async def initialize(self): + """Initialize the goal conversation service""" + logger.info("Initializing GoalConversationService") + + try: + self.pool = await asyncpg.create_pool( + self.database_url, min_size=1, max_size=5, command_timeout=60 + ) + + logger.info("GoalConversationService initialized successfully") + + except Exception as e: + logger.error(f"Failed to initialize GoalConversationService: {e}") + raise + + async def close(self): + """Close database connections""" + if self.pool: + await self.pool.close() + logger.info("GoalConversationService closed") + + async def create_goal_conversation( + self, + goal_id: str, + conversation_type: ConversationType, + conversation_title: str, + initial_context: Optional[Dict[str, Any]] = None, + participants: Optional[List[Dict[str, Any]]] = None, + created_by: Optional[str] = None, + ) -> str: + """Create a new conversation for a goal""" + + conversation_id = str(uuid.uuid4()) + + if initial_context is None: + initial_context = {} + if participants is None: + participants = [] + + try: + async with self.pool.acquire() as conn: + # Get goal context + goal = await conn.fetchrow( + """ + SELECT title, description, goal_type, target_deadline, + progress_percentage, current_value, target_value + FROM organization_goals WHERE id = $1 + """, + goal_id, + ) + + if not goal: + raise ValueError(f"Goal {goal_id} not found") + + # Enhanced context with goal information + enhanced_context = { + **initial_context, + "goal_title": goal["title"], + "goal_type": goal["goal_type"], + "goal_progress": float(goal["progress_percentage"]), + "days_to_deadline": ( + goal["target_deadline"] - datetime.now().date() + ).days, + "conversation_created_at": datetime.now().isoformat(), + } + + # Create conversation + await conn.execute( + """ + INSERT INTO goal_conversations ( + id, goal_id, conversation_type, conversation_title, + conversation_context, participants, status, created_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + """, + conversation_id, + goal_id, + conversation_type.value, + conversation_title, + json.dumps(enhanced_context), + json.dumps(participants), + ConversationStatus.ACTIVE.value, + created_by, + ) + + # Add initial system message with conversation starter + starter_message = self._generate_conversation_starter( + conversation_type, goal, enhanced_context + ) + + await self._add_message( + conversation_id, + MessageType.SYSTEM, + None, + "System", + starter_message, + {"conversation_starter": True}, + ) + + self.conversations_created += 1 + logger.info(f"Created conversation {conversation_id} for goal {goal_id}") + + return conversation_id + + except Exception as e: + logger.error(f"Error creating conversation: {e}") + raise + + async def add_message_to_conversation( + self, + conversation_id: str, + message_type: MessageType, + sender_id: Optional[str], + sender_name: str, + content: str, + metadata: Optional[Dict[str, Any]] = None, + references: Optional[List[str]] = None, + ) -> str: + """Add a message to a conversation""" + + if metadata is None: + metadata = {} + if references is None: + references = [] + + try: + # Add the message + message_id = await self._add_message( + conversation_id, + message_type, + sender_id, + sender_name, + content, + metadata, + references, + ) + + # Trigger conversation analysis for insights + await self._analyze_conversation_for_insights(conversation_id) + + # Update conversation activity timestamp + async with self.pool.acquire() as conn: + await conn.execute( + """ + UPDATE goal_conversations + SET last_activity_at = NOW(), updated_at = NOW() + WHERE id = $1 + """, + conversation_id, + ) + + self.messages_processed += 1 + + return message_id + + except Exception as e: + logger.error(f"Error adding message to conversation {conversation_id}: {e}") + raise + + async def get_conversation(self, conversation_id: str) -> Optional[Dict[str, Any]]: + """Get full conversation with messages, insights, and action items""" + + try: + async with self.pool.acquire() as conn: + # Get conversation details + conversation = await conn.fetchrow( + """ + SELECT gc.*, og.title as goal_title, og.goal_type + FROM goal_conversations gc + JOIN organization_goals og ON gc.goal_id = og.id + WHERE gc.id = $1 + """, + conversation_id, + ) + + if not conversation: + return None + + # Get messages + messages = ( + json.loads(conversation["messages"]) + if conversation["messages"] + else [] + ) + + # Get insights + insights = ( + json.loads(conversation["insights_generated"]) + if conversation["insights_generated"] + else [] + ) + + # Get action items + action_items = ( + json.loads(conversation["action_items"]) + if conversation["action_items"] + else [] + ) + + return { + "id": str(conversation["id"]), + "goal_id": str(conversation["goal_id"]), + "goal_title": conversation["goal_title"], + "conversation_type": conversation["conversation_type"], + "conversation_title": conversation["conversation_title"], + "conversation_summary": conversation["conversation_summary"], + "conversation_context": ( + json.loads(conversation["conversation_context"]) + if conversation["conversation_context"] + else {} + ), + "participants": ( + json.loads(conversation["participants"]) + if conversation["participants"] + else [] + ), + "messages": messages, + "insights_generated": insights, + "action_items": action_items, + "status": conversation["status"], + "last_activity_at": ( + conversation["last_activity_at"].isoformat() + if conversation["last_activity_at"] + else None + ), + "created_at": conversation["created_at"].isoformat(), + "updated_at": conversation["updated_at"].isoformat(), + "message_count": len(messages), + "insight_count": len(insights), + "action_item_count": len(action_items), + } + + except Exception as e: + logger.error(f"Error getting conversation {conversation_id}: {e}") + return None + + async def generate_planning_milestones( + self, conversation_id: str, planning_context: Optional[Dict[str, Any]] = None + ) -> List[Dict[str, Any]]: + """Generate milestone suggestions based on conversation analysis""" + + try: + async with self.pool.acquire() as conn: + # Get conversation and goal context + conversation = await conn.fetchrow( + """ + SELECT gc.*, og.title, og.description, og.goal_type, + og.target_deadline, og.target_value, og.target_unit + FROM goal_conversations gc + JOIN organization_goals og ON gc.goal_id = og.id + WHERE gc.id = $1 + """, + conversation_id, + ) + + if not conversation: + raise ValueError(f"Conversation {conversation_id} not found") + + # Analyze conversation content for milestone ideas + messages = ( + json.loads(conversation["messages"]) + if conversation["messages"] + else [] + ) + milestone_suggestions = self._extract_milestone_ideas_from_conversation( + messages, conversation, planning_context + ) + + # Generate AI-powered milestone recommendations + ai_milestones = await self._generate_ai_milestone_recommendations( + conversation, milestone_suggestions, planning_context + ) + + # Add milestones as insights to the conversation + milestone_insight = { + "id": str(uuid.uuid4()), + "insight_type": "milestone_recommendations", + "title": "AI-Generated Milestone Recommendations", + "description": f"Based on conversation analysis, here are {len(ai_milestones)} recommended milestones", + "confidence_score": 0.85, + "supporting_messages": [ + msg["id"] for msg in messages[-5:] if "id" in msg + ], # Last 5 messages + "suggested_actions": [ + { + "action": "create_milestones", + "description": "Create these milestones for the goal", + "milestones": ai_milestones, + } + ], + "generated_at": datetime.now().isoformat(), + } + + # Update conversation with milestone insight + await self._add_insight_to_conversation( + conversation_id, milestone_insight + ) + + return ai_milestones + + except Exception as e: + logger.error(f"Error generating planning milestones: {e}") + return [] + + async def conduct_progress_review( + self, conversation_id: str, review_period_days: int = 30 + ) -> Dict[str, Any]: + """Conduct AI-powered progress review for a goal conversation""" + + try: + async with self.pool.acquire() as conn: + # Get conversation and goal data + conversation = await conn.fetchrow( + """ + SELECT gc.*, og.* + FROM goal_conversations gc + JOIN organization_goals og ON gc.goal_id = og.id + WHERE gc.id = $1 + """, + conversation_id, + ) + + if not conversation: + raise ValueError(f"Conversation {conversation_id} not found") + + # Get recent progress data + progress_data = await conn.fetch( + """ + SELECT * FROM goal_progress_tracking + WHERE goal_id = $1 + AND recorded_at >= NOW() - INTERVAL '%s days' + ORDER BY recorded_at DESC + """, + str(conversation["goal_id"]), + review_period_days, + ) + + # Get milestones and tasks status + milestone_status = await conn.fetch( + """ + SELECT status, COUNT(*) as count + FROM goal_milestones + WHERE goal_id = $1 + GROUP BY status + """, + str(conversation["goal_id"]), + ) + + task_status = await conn.fetch( + """ + SELECT status, COUNT(*) as count + FROM goal_tasks + WHERE goal_id = $1 + GROUP BY status + """, + str(conversation["goal_id"]), + ) + + # Generate review analysis + review_analysis = { + "review_period_days": review_period_days, + "goal_progress": { + "current_progress": float(conversation["progress_percentage"]), + "target_value": ( + float(conversation["target_value"]) + if conversation["target_value"] + else None + ), + "current_value": ( + float(conversation["current_value"]) + if conversation["current_value"] + else None + ), + "completion_confidence": float( + conversation["completion_confidence"] + ), + }, + "milestone_summary": { + row["status"]: row["count"] for row in milestone_status + }, + "task_summary": { + row["status"]: row["count"] for row in task_status + }, + "progress_trend": self._calculate_progress_trend( + [dict(p) for p in progress_data] + ), + "risk_assessment": self._assess_goal_risks( + conversation, progress_data + ), + "recommendations": self._generate_progress_recommendations( + conversation, progress_data + ), + } + + # Add review as a structured message + review_message = self._format_progress_review_message(review_analysis) + await self._add_message( + conversation_id, + MessageType.AI_ANALYSIS, + None, + "Progress Analyzer", + review_message, + {"review_analysis": review_analysis}, + ) + + return review_analysis + + except Exception as e: + logger.error(f"Error conducting progress review: {e}") + return {"error": str(e)} + + async def extract_action_items_from_conversation( + self, conversation_id: str, auto_assign: bool = True + ) -> List[Dict[str, Any]]: + """Extract and create action items from conversation analysis""" + + try: + async with self.pool.acquire() as conn: + conversation = await conn.fetchrow( + """ + SELECT * FROM goal_conversations WHERE id = $1 + """, + conversation_id, + ) + + if not conversation: + raise ValueError(f"Conversation {conversation_id} not found") + + messages = ( + json.loads(conversation["messages"]) + if conversation["messages"] + else [] + ) + + # Analyze messages for actionable items + potential_actions = self._identify_action_items_in_messages(messages) + + # Convert to action item format + action_items = [] + for action in potential_actions: + if action["confidence"] >= self.auto_action_item_threshold: + action_item = { + "id": str(uuid.uuid4()), + "title": action["title"], + "description": action["description"], + "assigned_to": ( + action.get("assigned_to") if auto_assign else None + ), + "due_date": action.get("due_date"), + "status": "pending", + "priority": action.get("priority", 5), + "source_messages": action["source_messages"], + "created_at": datetime.now().isoformat(), + "confidence_score": action["confidence"], + } + action_items.append(action_item) + + # Update conversation with action items + if action_items: + existing_actions = ( + json.loads(conversation["action_items"]) + if conversation["action_items"] + else [] + ) + all_actions = existing_actions + action_items + + await conn.execute( + """ + UPDATE goal_conversations + SET action_items = $2, updated_at = NOW() + WHERE id = $1 + """, + conversation_id, + json.dumps(all_actions), + ) + + self.action_items_created += len(action_items) + + return action_items + + except Exception as e: + logger.error(f"Error extracting action items: {e}") + return [] + + async def get_goal_conversations( + self, + goal_id: str, + conversation_type: Optional[ConversationType] = None, + status: Optional[ConversationStatus] = None, + limit: int = 10, + ) -> List[Dict[str, Any]]: + """Get conversations for a goal with optional filtering""" + + try: + async with self.pool.acquire() as conn: + where_conditions = ["goal_id = $1"] + params = [goal_id] + param_idx = 2 + + if conversation_type: + where_conditions.append(f"conversation_type = ${param_idx}") + params.append(conversation_type.value) + param_idx += 1 + + if status: + where_conditions.append(f"status = ${param_idx}") + params.append(status.value) + param_idx += 1 + + where_clause = " AND ".join(where_conditions) + + conversations = await conn.fetch( + """ + SELECT id, conversation_type, conversation_title, conversation_summary, + status, last_activity_at, created_at, + COALESCE(array_length(string_to_array(messages::text, '}}'), 1), 0) as message_count, + COALESCE(array_length(string_to_array(action_items::text, '}}'), 1), 0) as action_count + FROM goal_conversations + WHERE {where_clause} + ORDER BY last_activity_at DESC, created_at DESC + LIMIT ${param_idx} + """.format( # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params + where_clause=where_clause, param_idx=param_idx + ), + *params, + limit, + ) + + return [dict(conv) for conv in conversations] + + except Exception as e: + logger.error(f"Error getting conversations for goal {goal_id}: {e}") + return [] + + # Helper methods for conversation management + + async def _add_message( + self, + conversation_id: str, + message_type: MessageType, + sender_id: Optional[str], + sender_name: str, + content: str, + metadata: Dict[str, Any], + references: Optional[List[str]] = None, + ) -> str: + """Add a message to conversation""" + + message_id = str(uuid.uuid4()) + message = { + "id": message_id, + "message_type": message_type.value, + "sender_id": sender_id, + "sender_name": sender_name, + "content": content, + "metadata": metadata, + "timestamp": datetime.now().isoformat(), + "references": references or [], + "reactions": [], + } + + async with self.pool.acquire() as conn: + # Get current messages + current_messages = await conn.fetchval( + """ + SELECT messages FROM goal_conversations WHERE id = $1 + """, + conversation_id, + ) + + messages = json.loads(current_messages) if current_messages else [] + messages.append(message) + + # Limit message history + if len(messages) > self.max_conversation_messages: + messages = messages[-self.max_conversation_messages :] + + # Update conversation + await conn.execute( + """ + UPDATE goal_conversations + SET messages = $2, updated_at = NOW() + WHERE id = $1 + """, + conversation_id, + json.dumps(messages), + ) + + return message_id + + async def _add_insight_to_conversation( + self, conversation_id: str, insight: Dict[str, Any] + ): + """Add an insight to conversation""" + + async with self.pool.acquire() as conn: + # Get current insights + current_insights = await conn.fetchval( + """ + SELECT insights_generated FROM goal_conversations WHERE id = $1 + """, + conversation_id, + ) + + insights = json.loads(current_insights) if current_insights else [] + insights.append(insight) + + # Update conversation + await conn.execute( + """ + UPDATE goal_conversations + SET insights_generated = $2, updated_at = NOW() + WHERE id = $1 + """, + conversation_id, + json.dumps(insights), + ) + + self.insights_generated += 1 + + def _generate_conversation_starter( + self, + conversation_type: ConversationType, + goal: Dict[str, Any], + context: Dict[str, Any], + ) -> str: + """Generate an appropriate conversation starter""" + + starters = self.conversation_starters.get(conversation_type, {}) + + if conversation_type == ConversationType.PLANNING: + return starters["default"].format( + goal_title=goal["title"], + goal_type=goal["goal_type"], + deadline=goal["target_deadline"].strftime("%B %d, %Y"), + target_value=( + goal["target_value"] + if goal["target_value"] + else "defined objectives" + ), + ) + elif conversation_type == ConversationType.REVIEW: + return starters["default"].format( + goal_title=goal["title"], + current_progress=f"{goal['progress_percentage']:.1f}%", + ) + else: + return starters.get("default", f"Let's discuss the {goal['title']} goal.") + + def _initialize_conversation_starters(self) -> Dict[str, Dict[str, str]]: + """Initialize conversation starter templates""" + + return { + ConversationType.PLANNING: { + "default": """Welcome to the strategic planning session for "{goal_title}"! + +🎯 **Goal**: {goal_title} +📊 **Type**: {goal_type} +📅 **Deadline**: {deadline} +🎌 **Target**: {target_value} + +Let's break this goal down into actionable milestones and tasks. Here are some questions to get us started: + +1. **What are the major milestones we need to achieve?** +2. **What dependencies and blockers should we consider?** +3. **Which teams and resources will be involved?** +4. **How should we measure progress along the way?** + +What aspect would you like to focus on first?""" + }, + ConversationType.REVIEW: { + "default": """Time for a progress review of "{goal_title}"! + +📈 **Current Progress**: {current_progress} + +Let's evaluate our progress, identify what's working well, and address any challenges. + +Key areas to discuss: +- Recent achievements and wins +- Current blockers or risks +- Resource allocation and team performance +- Timeline adjustments if needed +- Next steps and priorities + +What would you like to review first?""" + }, + ConversationType.PROBLEM_SOLVING: { + "default": """Problem-solving session for "{goal_title}". + +Let's identify the specific challenges we're facing and work together to find solutions. Please share: +- What specific problems or blockers have emerged? +- What have we tried so far? +- What constraints or requirements should we consider? + +What's the main challenge you'd like to tackle?""" + }, + } + + def _initialize_analysis_prompts(self) -> Dict[str, str]: + """Initialize AI analysis prompts for conversation processing""" + + return { + "extract_milestones": """Analyze this goal conversation and extract potential milestones mentioned or implied. Look for: +- Time-based deliverables or checkpoints +- Measurable objectives or targets +- Dependencies between activities +- Key decision points or reviews + +Return milestones with titles, descriptions, and target dates.""", + "identify_risks": """Analyze this conversation for potential risks, blockers, or concerns mentioned. Look for: +- Resource constraints or availability issues +- Technical challenges or unknowns +- Timeline concerns or dependencies +- Team capacity or skill gaps +- External dependencies or market factors + +Assess the probability and impact of each risk.""", + "extract_actions": """Extract specific action items from this conversation. Look for: +- Tasks or activities that someone needs to do +- Decisions that need to be made +- Information that needs to be gathered +- People who need to be contacted +- Deadlines or time-sensitive items + +Include who should be responsible and when it should be done.""", + } + + # Placeholder implementations for AI analysis methods + + async def _analyze_conversation_for_insights(self, conversation_id: str): + """Analyze conversation for insights (placeholder for AI integration)""" + # This would integrate with an AI service for conversation analysis + pass + + def _extract_milestone_ideas_from_conversation( + self, + messages: List[Dict[str, Any]], + conversation: Dict[str, Any], + planning_context: Optional[Dict[str, Any]], + ) -> List[Dict[str, Any]]: + """Extract milestone ideas from conversation messages""" + # Simplified implementation - would use NLP/AI for real extraction + milestone_keywords = [ + "milestone", + "phase", + "deliverable", + "target", + "deadline", + "complete", + ] + + milestones = [] + for message in messages: + content = message.get("content", "").lower() + if any(keyword in content for keyword in milestone_keywords): + # Extract potential milestone (simplified) + milestones.append( + { + "title": f"Milestone from conversation", + "description": message.get("content", "")[:200], + "source_message": message.get("id"), + "confidence": 0.7, + } + ) + + return milestones[:5] # Return top 5 candidates + + async def _generate_ai_milestone_recommendations( + self, + conversation: Dict[str, Any], + milestone_suggestions: List[Dict[str, Any]], + planning_context: Optional[Dict[str, Any]], + ) -> List[Dict[str, Any]]: + """Generate AI-powered milestone recommendations""" + # Placeholder implementation - would use AI for intelligent milestone generation + + goal_type = conversation["goal_type"] + timeline_days = (conversation["target_deadline"] - datetime.now().date()).days + + # Generate sample milestones based on goal type + if goal_type == "business" and "revenue" in conversation["title"].lower(): + return [ + { + "title": "Foundation Setup", + "description": "Establish initial infrastructure, team structure, and processes", + "target_date": (datetime.now() + timedelta(days=timeline_days // 4)) + .date() + .isoformat(), + "milestone_type": "checkpoint", + }, + { + "title": "Growth Phase Launch", + "description": "Execute marketing campaigns and sales initiatives", + "target_date": (datetime.now() + timedelta(days=timeline_days // 2)) + .date() + .isoformat(), + "milestone_type": "deliverable", + }, + { + "title": "Scale and Optimize", + "description": "Optimize processes and scale operations for target achievement", + "target_date": ( + datetime.now() + timedelta(days=timeline_days * 3 // 4) + ) + .date() + .isoformat(), + "milestone_type": "metric", + }, + ] + + return [] + + # Additional helper methods for conversation analysis + def _calculate_progress_trend(self, progress_data: List[Dict[str, Any]]) -> str: + """Calculate progress trend from historical data""" + if len(progress_data) < 2: + return "insufficient_data" + + # Simple trend calculation + recent_progress = progress_data[0]["progress_percentage"] + older_progress = progress_data[-1]["progress_percentage"] + + if recent_progress > older_progress * 1.1: + return "accelerating" + elif recent_progress < older_progress * 0.9: + return "declining" + else: + return "steady" + + def _assess_goal_risks( + self, goal: Dict[str, Any], progress_data: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + """Assess risks based on goal and progress data""" + risks = [] + + # Timeline risk + days_remaining = (goal["target_deadline"] - datetime.now().date()).days + progress = float(goal["progress_percentage"]) + + if days_remaining < 30 and progress < 70: + risks.append( + { + "type": "timeline_risk", + "severity": "high", + "description": "Goal progress is behind schedule with limited time remaining", + } + ) + + return risks + + def _generate_progress_recommendations( + self, goal: Dict[str, Any], progress_data: List[Dict[str, Any]] + ) -> List[str]: + """Generate recommendations based on progress analysis""" + recommendations = [] + + progress = float(goal["progress_percentage"]) + + if progress < 25: + recommendations.append( + "Consider breaking down remaining work into smaller, more manageable tasks" + ) + + if progress > 75: + recommendations.append( + "Focus on final quality checks and prepare for goal completion" + ) + + return recommendations + + def _identify_action_items_in_messages( + self, messages: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + """Identify potential action items from messages""" + # Simplified implementation - would use NLP for real extraction + action_keywords = [ + "need to", + "should", + "must", + "will", + "action", + "task", + "todo", + ] + + actions = [] + for message in messages: + content = message.get("content", "").lower() + if any(keyword in content for keyword in action_keywords): + actions.append( + { + "title": f"Action item from conversation", + "description": message.get("content", "")[:200], + "source_messages": [message.get("id")], + "confidence": 0.8, + "priority": 5, + } + ) + + return actions[:10] # Return top 10 candidates + + def _format_progress_review_message(self, review_analysis: Dict[str, Any]) -> str: + """Format progress review analysis as a readable message""" + + progress = review_analysis["goal_progress"]["current_progress"] + trend = review_analysis["progress_trend"] + + message = f"""## Progress Review Summary + +**Current Progress**: {progress:.1f}% +**Trend**: {trend.replace('_', ' ').title()} + +### Key Metrics +""" + + if review_analysis["milestone_summary"]: + message += "\n**Milestones:**\n" + for status, count in review_analysis["milestone_summary"].items(): + message += f"- {status.replace('_', ' ').title()}: {count}\n" + + if review_analysis["recommendations"]: + message += "\n### Recommendations\n" + for i, rec in enumerate(review_analysis["recommendations"], 1): + message += f"{i}. {rec}\n" + + return message diff --git a/services/orchestrator/hierarchy_endpoints.py b/services/orchestrator/hierarchy_endpoints.py index 95d2d86..083d97f 100644 --- a/services/orchestrator/hierarchy_endpoints.py +++ b/services/orchestrator/hierarchy_endpoints.py @@ -1,393 +1,393 @@ -import asyncio -import json -from typing import Any, Dict, List, Optional - -import httpx -from fastapi import APIRouter, HTTPException - -from database import DatabaseManager - -router = APIRouter(prefix="/hierarchy", tags=["hierarchy"]) - - -@router.get("/visualization") -async def get_hierarchy_visualization(): - """ - Get complete organizational hierarchy for visualization with ReactFlow/GoJS - - Returns structured data optimized for hierarchical visualization libraries: - - Organizations as root nodes - - Teams as intermediate nodes - - Agents as leaf nodes - - Relationships and positioning data - """ - try: - # Get data from hierarchy API - async with httpx.AsyncClient() as client: - # Get all organizations - orgs_response = await client.get("http://localhost:8006/organizations") - organizations = orgs_response.json() - - # Get all teams - teams_response = await client.get("http://localhost:8006/teams") - all_teams = teams_response.json() - - # Get all agents - agents_response = await client.get("http://localhost:8006/agents") - all_agents = agents_response.json() - - if not organizations: - return {"nodes": [], "edges": [], "message": "No organizations found"} - - nodes = [] - edges = [] - y_offset = 0 - - for org_idx, org in enumerate(organizations): - org_id = org["id"] - org_node_id = f"org-{org_id}" - - # Add organization node - nodes.append( - { - "id": org_node_id, - "type": "organization", - "data": { - "label": org["name"], - "description": org.get("description", ""), - "type": "Organization", - "settings": org.get("settings", {}), - "entity_id": org_id, - }, - "position": {"x": org_idx * 800, "y": y_offset}, - "style": { - "background": "#1e40af", - "color": "white", - "border": "2px solid #1e3a8a", - "borderRadius": "12px", - "padding": "12px", - "minWidth": "200px", - }, - } - ) - - # Filter teams for this organization - teams = [t for t in all_teams if t.get("organization_id") == org_id] - team_y_offset = y_offset + 150 - - for team_idx, team in enumerate(teams): - team_id = team["id"] - team_node_id = f"team-{team_id}" - - # Add team node - nodes.append( - { - "id": team_node_id, - "type": "team", - "data": { - "label": team["name"], - "description": team.get("description", ""), - "type": f"Team ({team.get('team_type', 'general')})", - "settings": team.get("settings", {}), - "entity_id": team_id, - "organization_id": org_id, - }, - "position": { - "x": org_idx * 800 + (team_idx % 3) * 250 - 250, - "y": team_y_offset + (team_idx // 3) * 120, - }, - "style": { - "background": "#059669", - "color": "white", - "border": "2px solid #047857", - "borderRadius": "8px", - "padding": "10px", - "minWidth": "180px", - }, - } - ) - - # Add edge from organization to team - edges.append( - { - "id": f"edge-{org_node_id}-{team_node_id}", - "source": org_node_id, - "target": team_node_id, - "type": "smoothstep", - "style": {"stroke": "#64748b", "strokeWidth": 2}, - "markerEnd": {"type": "arrowclosed", "color": "#64748b"}, - } - ) - - # Filter agents for this team - note: need to check team_id in agents - agents = [ - a - for a in all_agents - if a.get("team_id") == team_id - or (hasattr(a, "config") and a.config.get("team_id") == team_id) - ] - # Fallback: if no team_id in agent data, we'll have limited agents - if not agents and all_agents: - # For now, include some agents if they don't have team assignments - agents = all_agents[:2] # Include first 2 agents as examples - - agent_y_offset = team_y_offset + (team_idx // 3) * 120 + 100 - - for agent_idx, agent in enumerate(agents): - agent_id = agent["id"] - agent_node_id = f"agent-{agent_id}" - - # Determine agent color by type - agent_colors = { - "executive": {"bg": "#dc2626", "border": "#b91c1c"}, - "developer": {"bg": "#2563eb", "border": "#1d4ed8"}, - "marketing": {"bg": "#7c3aed", "border": "#6d28d9"}, - "sales": {"bg": "#ea580c", "border": "#c2410c"}, - "qa": {"bg": "#16a34a", "border": "#15803d"}, - "devops": {"bg": "#0891b2", "border": "#0e7490"}, - "designer": {"bg": "#e11d48", "border": "#be185d"}, - } - - agent_type = agent.get("type", "developer") - colors = agent_colors.get( - agent_type, {"bg": "#6b7280", "border": "#4b5563"} - ) - - # Add agent node - nodes.append( - { - "id": agent_node_id, - "type": "agent", - "data": { - "label": agent["name"], - "role": agent.get("role", ""), - "type": f"Agent ({agent_type})", - "status": agent.get("status", "active"), - "config": agent.get("config", {}), - "entity_id": agent_id, - "team_id": team_id, - "organization_id": org_id, - }, - "position": { - "x": org_idx * 800 - + (team_idx % 3) * 250 - - 250 - + (agent_idx % 2) * 120 - - 60, - "y": agent_y_offset + (agent_idx // 2) * 80, - }, - "style": { - "background": colors["bg"], - "color": "white", - "border": f"2px solid {colors['border']}", - "borderRadius": "6px", - "padding": "8px", - "minWidth": "150px", - }, - } - ) - - # Add edge from team to agent - edges.append( - { - "id": f"edge-{team_node_id}-{agent_node_id}", - "source": team_node_id, - "target": agent_node_id, - "type": "smoothstep", - "style": {"stroke": "#94a3b8", "strokeWidth": 1.5}, - "markerEnd": {"type": "arrowclosed", "color": "#94a3b8"}, - } - ) - - return { - "nodes": nodes, - "edges": edges, - "metadata": { - "total_organizations": len(organizations), - "total_teams": len(all_teams), - "total_agents": len(all_agents), - "generated_at": "2025-08-06T10:00:00Z", - }, - } - - except Exception as e: - raise HTTPException( - status_code=500, - detail=f"Failed to generate hierarchy visualization: {str(e)}", - ) - - -@router.get("/stats") -async def get_hierarchy_stats(): - """Get comprehensive hierarchy statistics""" - try: - organizations = await DatabaseManager.get_organizations() - - if not organizations: - return {"organizations": 0, "teams": 0, "agents": 0, "by_organization": []} - - stats = { - "organizations": len(organizations), - "teams": 0, - "agents": 0, - "by_organization": [], - "agent_types": {}, - "team_types": {}, - } - - for org in organizations: - org_id = org["id"] - teams = await DatabaseManager.get_teams(org_id) - org_agent_count = 0 - org_teams_by_type = {} - org_agents_by_type = {} - - for team in teams: - team_id = team["id"] - team_type = team.get("team_type", "general") - org_teams_by_type[team_type] = org_teams_by_type.get(team_type, 0) + 1 - stats["team_types"][team_type] = ( - stats["team_types"].get(team_type, 0) + 1 - ) - - agents = await DatabaseManager.get_agents(team_id) - org_agent_count += len(agents) - - for agent in agents: - agent_type = agent.get("type", "developer") - org_agents_by_type[agent_type] = ( - org_agents_by_type.get(agent_type, 0) + 1 - ) - stats["agent_types"][agent_type] = ( - stats["agent_types"].get(agent_type, 0) + 1 - ) - - stats["by_organization"].append( - { - "id": org_id, - "name": org["name"], - "teams": len(teams), - "agents": org_agent_count, - "teams_by_type": org_teams_by_type, - "agents_by_type": org_agents_by_type, - } - ) - - stats["teams"] += len(teams) - stats["agents"] += org_agent_count - - return stats - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get hierarchy stats: {str(e)}" - ) - - -@router.get("/organization/{organization_id}/chart") -async def get_organization_chart(organization_id: str): - """Get detailed chart data for a specific organization""" - try: - # Get organization details - organization = await DatabaseManager.get_organization(organization_id) - if not organization: - raise HTTPException(status_code=404, detail="Organization not found") - - # Get teams and agents - teams = await DatabaseManager.get_teams(organization_id) - - chart_data = {"organization": organization, "teams": [], "total_agents": 0} - - for team in teams: - team_id = team["id"] - agents = await DatabaseManager.get_agents(team_id) - - team_data = { - "id": team_id, - "name": team["name"], - "description": team.get("description", ""), - "team_type": team.get("team_type", "general"), - "settings": team.get("settings", {}), - "agents": agents, - "agent_count": len(agents), - } - - chart_data["teams"].append(team_data) - chart_data["total_agents"] += len(agents) - - return chart_data - - except HTTPException: - raise - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get organization chart: {str(e)}" - ) - - -@router.get("/search") -async def search_hierarchy(q: str, entity_type: Optional[str] = None): - """Search across organizations, teams, and agents""" - try: - if not q or len(q) < 2: - raise HTTPException( - status_code=400, detail="Query must be at least 2 characters" - ) - - results = {"organizations": [], "teams": [], "agents": [], "total_results": 0} - - query = q.lower() - - # Search organizations - if not entity_type or entity_type == "organization": - organizations = await DatabaseManager.get_organizations() - for org in organizations: - if ( - query in org["name"].lower() - or query in org.get("description", "").lower() - ): - results["organizations"].append(org) - - # Search teams - if not entity_type or entity_type == "team": - organizations = await DatabaseManager.get_organizations() - for org in organizations: - teams = await DatabaseManager.get_teams(org["id"]) - for team in teams: - if ( - query in team["name"].lower() - or query in team.get("description", "").lower() - or query in team.get("team_type", "").lower() - ): - team["organization_name"] = org["name"] - results["teams"].append(team) - - # Search agents - if not entity_type or entity_type == "agent": - organizations = await DatabaseManager.get_organizations() - for org in organizations: - teams = await DatabaseManager.get_teams(org["id"]) - for team in teams: - agents = await DatabaseManager.get_agents(team["id"]) - for agent in agents: - if ( - query in agent["name"].lower() - or query in agent.get("role", "").lower() - or query in agent.get("type", "").lower() - ): - agent["organization_name"] = org["name"] - agent["team_name"] = team["name"] - results["agents"].append(agent) - - results["total_results"] = ( - len(results["organizations"]) - + len(results["teams"]) - + len(results["agents"]) - ) - - return results - - except HTTPException: - raise - except Exception as e: - raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}") +import asyncio +import json +from typing import Any, Dict, List, Optional + +import httpx +from fastapi import APIRouter, HTTPException + +from database import DatabaseManager + +router = APIRouter(prefix="/hierarchy", tags=["hierarchy"]) + + +@router.get("/visualization") +async def get_hierarchy_visualization(): + """ + Get complete organizational hierarchy for visualization with ReactFlow/GoJS + + Returns structured data optimized for hierarchical visualization libraries: + - Organizations as root nodes + - Teams as intermediate nodes + - Agents as leaf nodes + - Relationships and positioning data + """ + try: + # Get data from hierarchy API + async with httpx.AsyncClient() as client: + # Get all organizations + orgs_response = await client.get("http://localhost:8006/organizations") + organizations = orgs_response.json() + + # Get all teams + teams_response = await client.get("http://localhost:8006/teams") + all_teams = teams_response.json() + + # Get all agents + agents_response = await client.get("http://localhost:8006/agents") + all_agents = agents_response.json() + + if not organizations: + return {"nodes": [], "edges": [], "message": "No organizations found"} + + nodes = [] + edges = [] + y_offset = 0 + + for org_idx, org in enumerate(organizations): + org_id = org["id"] + org_node_id = f"org-{org_id}" + + # Add organization node + nodes.append( + { + "id": org_node_id, + "type": "organization", + "data": { + "label": org["name"], + "description": org.get("description", ""), + "type": "Organization", + "settings": org.get("settings", {}), + "entity_id": org_id, + }, + "position": {"x": org_idx * 800, "y": y_offset}, + "style": { + "background": "#1e40af", + "color": "white", + "border": "2px solid #1e3a8a", + "borderRadius": "12px", + "padding": "12px", + "minWidth": "200px", + }, + } + ) + + # Filter teams for this organization + teams = [t for t in all_teams if t.get("organization_id") == org_id] + team_y_offset = y_offset + 150 + + for team_idx, team in enumerate(teams): + team_id = team["id"] + team_node_id = f"team-{team_id}" + + # Add team node + nodes.append( + { + "id": team_node_id, + "type": "team", + "data": { + "label": team["name"], + "description": team.get("description", ""), + "type": f"Team ({team.get('team_type', 'general')})", + "settings": team.get("settings", {}), + "entity_id": team_id, + "organization_id": org_id, + }, + "position": { + "x": org_idx * 800 + (team_idx % 3) * 250 - 250, + "y": team_y_offset + (team_idx // 3) * 120, + }, + "style": { + "background": "#059669", + "color": "white", + "border": "2px solid #047857", + "borderRadius": "8px", + "padding": "10px", + "minWidth": "180px", + }, + } + ) + + # Add edge from organization to team + edges.append( + { + "id": f"edge-{org_node_id}-{team_node_id}", + "source": org_node_id, + "target": team_node_id, + "type": "smoothstep", + "style": {"stroke": "#64748b", "strokeWidth": 2}, + "markerEnd": {"type": "arrowclosed", "color": "#64748b"}, + } + ) + + # Filter agents for this team - note: need to check team_id in agents + agents = [ + a + for a in all_agents + if a.get("team_id") == team_id + or (hasattr(a, "config") and a.config.get("team_id") == team_id) + ] + # Fallback: if no team_id in agent data, we'll have limited agents + if not agents and all_agents: + # For now, include some agents if they don't have team assignments + agents = all_agents[:2] # Include first 2 agents as examples + + agent_y_offset = team_y_offset + (team_idx // 3) * 120 + 100 + + for agent_idx, agent in enumerate(agents): + agent_id = agent["id"] + agent_node_id = f"agent-{agent_id}" + + # Determine agent color by type + agent_colors = { + "executive": {"bg": "#dc2626", "border": "#b91c1c"}, + "developer": {"bg": "#2563eb", "border": "#1d4ed8"}, + "marketing": {"bg": "#7c3aed", "border": "#6d28d9"}, + "sales": {"bg": "#ea580c", "border": "#c2410c"}, + "qa": {"bg": "#16a34a", "border": "#15803d"}, + "devops": {"bg": "#0891b2", "border": "#0e7490"}, + "designer": {"bg": "#e11d48", "border": "#be185d"}, + } + + agent_type = agent.get("type", "developer") + colors = agent_colors.get( + agent_type, {"bg": "#6b7280", "border": "#4b5563"} + ) + + # Add agent node + nodes.append( + { + "id": agent_node_id, + "type": "agent", + "data": { + "label": agent["name"], + "role": agent.get("role", ""), + "type": f"Agent ({agent_type})", + "status": agent.get("status", "active"), + "config": agent.get("config", {}), + "entity_id": agent_id, + "team_id": team_id, + "organization_id": org_id, + }, + "position": { + "x": org_idx * 800 + + (team_idx % 3) * 250 + - 250 + + (agent_idx % 2) * 120 + - 60, + "y": agent_y_offset + (agent_idx // 2) * 80, + }, + "style": { + "background": colors["bg"], + "color": "white", + "border": f"2px solid {colors['border']}", + "borderRadius": "6px", + "padding": "8px", + "minWidth": "150px", + }, + } + ) + + # Add edge from team to agent + edges.append( + { + "id": f"edge-{team_node_id}-{agent_node_id}", + "source": team_node_id, + "target": agent_node_id, + "type": "smoothstep", + "style": {"stroke": "#94a3b8", "strokeWidth": 1.5}, + "markerEnd": {"type": "arrowclosed", "color": "#94a3b8"}, + } + ) + + return { + "nodes": nodes, + "edges": edges, + "metadata": { + "total_organizations": len(organizations), + "total_teams": len(all_teams), + "total_agents": len(all_agents), + "generated_at": "2025-08-06T10:00:00Z", + }, + } + + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"Failed to generate hierarchy visualization: {str(e)}", + ) + + +@router.get("/stats") +async def get_hierarchy_stats(): + """Get comprehensive hierarchy statistics""" + try: + organizations = await DatabaseManager.get_organizations() + + if not organizations: + return {"organizations": 0, "teams": 0, "agents": 0, "by_organization": []} + + stats = { + "organizations": len(organizations), + "teams": 0, + "agents": 0, + "by_organization": [], + "agent_types": {}, + "team_types": {}, + } + + for org in organizations: + org_id = org["id"] + teams = await DatabaseManager.get_teams(org_id) + org_agent_count = 0 + org_teams_by_type = {} + org_agents_by_type = {} + + for team in teams: + team_id = team["id"] + team_type = team.get("team_type", "general") + org_teams_by_type[team_type] = org_teams_by_type.get(team_type, 0) + 1 + stats["team_types"][team_type] = ( + stats["team_types"].get(team_type, 0) + 1 + ) + + agents = await DatabaseManager.get_agents(team_id) + org_agent_count += len(agents) + + for agent in agents: + agent_type = agent.get("type", "developer") + org_agents_by_type[agent_type] = ( + org_agents_by_type.get(agent_type, 0) + 1 + ) + stats["agent_types"][agent_type] = ( + stats["agent_types"].get(agent_type, 0) + 1 + ) + + stats["by_organization"].append( + { + "id": org_id, + "name": org["name"], + "teams": len(teams), + "agents": org_agent_count, + "teams_by_type": org_teams_by_type, + "agents_by_type": org_agents_by_type, + } + ) + + stats["teams"] += len(teams) + stats["agents"] += org_agent_count + + return stats + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get hierarchy stats: {str(e)}" + ) + + +@router.get("/organization/{organization_id}/chart") +async def get_organization_chart(organization_id: str): + """Get detailed chart data for a specific organization""" + try: + # Get organization details + organization = await DatabaseManager.get_organization(organization_id) + if not organization: + raise HTTPException(status_code=404, detail="Organization not found") + + # Get teams and agents + teams = await DatabaseManager.get_teams(organization_id) + + chart_data = {"organization": organization, "teams": [], "total_agents": 0} + + for team in teams: + team_id = team["id"] + agents = await DatabaseManager.get_agents(team_id) + + team_data = { + "id": team_id, + "name": team["name"], + "description": team.get("description", ""), + "team_type": team.get("team_type", "general"), + "settings": team.get("settings", {}), + "agents": agents, + "agent_count": len(agents), + } + + chart_data["teams"].append(team_data) + chart_data["total_agents"] += len(agents) + + return chart_data + + except HTTPException: + raise + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get organization chart: {str(e)}" + ) + + +@router.get("/search") +async def search_hierarchy(q: str, entity_type: Optional[str] = None): + """Search across organizations, teams, and agents""" + try: + if not q or len(q) < 2: + raise HTTPException( + status_code=400, detail="Query must be at least 2 characters" + ) + + results = {"organizations": [], "teams": [], "agents": [], "total_results": 0} + + query = q.lower() + + # Search organizations + if not entity_type or entity_type == "organization": + organizations = await DatabaseManager.get_organizations() + for org in organizations: + if ( + query in org["name"].lower() + or query in org.get("description", "").lower() + ): + results["organizations"].append(org) + + # Search teams + if not entity_type or entity_type == "team": + organizations = await DatabaseManager.get_organizations() + for org in organizations: + teams = await DatabaseManager.get_teams(org["id"]) + for team in teams: + if ( + query in team["name"].lower() + or query in team.get("description", "").lower() + or query in team.get("team_type", "").lower() + ): + team["organization_name"] = org["name"] + results["teams"].append(team) + + # Search agents + if not entity_type or entity_type == "agent": + organizations = await DatabaseManager.get_organizations() + for org in organizations: + teams = await DatabaseManager.get_teams(org["id"]) + for team in teams: + agents = await DatabaseManager.get_agents(team["id"]) + for agent in agents: + if ( + query in agent["name"].lower() + or query in agent.get("role", "").lower() + or query in agent.get("type", "").lower() + ): + agent["organization_name"] = org["name"] + agent["team_name"] = team["name"] + results["agents"].append(agent) + + results["total_results"] = ( + len(results["organizations"]) + + len(results["teams"]) + + len(results["agents"]) + ) + + return results + + except HTTPException: + raise + except Exception as e: + raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}") diff --git a/services/orchestrator/knowledge_propagation_engine.py b/services/orchestrator/knowledge_propagation_engine.py index 094b676..056b6f5 100644 --- a/services/orchestrator/knowledge_propagation_engine.py +++ b/services/orchestrator/knowledge_propagation_engine.py @@ -1,958 +1,958 @@ -""" -Knowledge Propagation Engine for FuzeAgent - -This module handles automated knowledge flow between agents, teams, and organizations. -It determines when knowledge should be propagated, executes the propagation, -and manages the lifecycle of knowledge across hierarchical levels. -""" - -import asyncio -import json -import logging -import uuid -from dataclasses import dataclass -from datetime import datetime, timedelta -from enum import Enum -from typing import Any, Dict, List, Optional, Set, Tuple - -import asyncpg -from sentence_transformers import SentenceTransformer - -from .organization_rag_manager import ( - ContentType, - KnowledgeCategory, - OrganizationRAGManager, - SourceType, - VisibilityLevel, -) -from .team_knowledge_manager import TeamKnowledgeManager - -logger = logging.getLogger(__name__) - - -class PropagationTrigger(str, Enum): - TASK_COMPLETION = "task_completion" - KNOWLEDGE_THRESHOLD = "knowledge_threshold" - MANUAL_REQUEST = "manual_request" - SCHEDULED_SYNC = "scheduled_sync" - CROSS_TEAM_REQUEST = "cross_team_request" - QUALITY_IMPROVEMENT = "quality_improvement" - - -class PropagationStatus(str, Enum): - PENDING = "pending" - PROCESSING = "processing" - COMPLETED = "completed" - FAILED = "failed" - REJECTED = "rejected" - - -class AcceptanceStatus(str, Enum): - PENDING = "pending" - ACCEPTED = "accepted" - REJECTED = "rejected" - MODIFIED = "modified" - - -@dataclass -class PropagationRule: - """Defines rules for knowledge propagation""" - - source_type: str # 'agent', 'team', 'organization' - target_type: str # 'agent', 'team', 'organization' - min_confidence: float - min_success_correlation: float - min_usage_count: int - knowledge_categories: List[KnowledgeCategory] - auto_approve: bool - propagation_weight: float - - -@dataclass -class PropagationTask: - """Represents a knowledge propagation task""" - - id: str - source_type: str - source_id: str - target_type: str - target_id: str - knowledge_type: str - knowledge_content_id: str - propagation_method: str - propagation_trigger: PropagationTrigger - confidence_score: float - propagation_status: PropagationStatus - acceptance_status: AcceptanceStatus - metadata: Dict[str, Any] - created_at: datetime - processed_at: Optional[datetime] - completed_at: Optional[datetime] - - -class KnowledgePropagationEngine: - """ - Manages the automated flow of knowledge across the organization hierarchy. - Handles agent → team → organization propagation and cross-team sharing. - """ - - def __init__( - self, - database_url: str, - org_rag_manager: OrganizationRAGManager, - team_knowledge_manager: TeamKnowledgeManager, - ): - self.database_url = database_url - self.org_rag_manager = org_rag_manager - self.team_knowledge_manager = team_knowledge_manager - self.pool: Optional[asyncpg.Pool] = None - - # Initialize embedding model for similarity analysis - self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") - - # Default propagation rules - self.default_rules = self._create_default_propagation_rules() - - # Configuration - self.propagation_batch_size = 50 - self.max_concurrent_propagations = 5 - self.similarity_threshold = 0.8 - self.propagation_cooldown_hours = 24 - - # Statistics - self.propagations_processed = 0 - self.propagations_completed = 0 - self.propagations_rejected = 0 - - # Background task management - self._propagation_task: Optional[asyncio.Task] = None - self._running = False - - async def initialize(self): - """Initialize the knowledge propagation engine""" - logger.info("Initializing KnowledgePropagationEngine") - - try: - self.pool = await asyncpg.create_pool( - self.database_url, min_size=2, max_size=10, command_timeout=60 - ) - - # Start background propagation processing - self._running = True - self._propagation_task = asyncio.create_task( - self._background_propagation_processor() - ) - - logger.info("KnowledgePropagationEngine initialized successfully") - - except Exception as e: - logger.error(f"Failed to initialize KnowledgePropagationEngine: {e}") - raise - - async def close(self): - """Close the propagation engine and cleanup resources""" - self._running = False - - if self._propagation_task: - self._propagation_task.cancel() - try: - await self._propagation_task - except asyncio.CancelledError: - pass - - if self.pool: - await self.pool.close() - - logger.info("KnowledgePropagationEngine closed") - - async def trigger_agent_to_team_propagation( - self, agent_id: str, task_id: str, task_outcome: Dict[str, Any] - ) -> List[str]: - """Trigger knowledge propagation from agent to team level after task completion""" - - propagation_ids = [] - - async with self.pool.acquire() as conn: - # Get agent's team - team_id = await conn.fetchval( - """ - SELECT team_id FROM agents WHERE id = $1 - """, - agent_id, - ) - - if not team_id: - logger.warning(f"No team found for agent {agent_id}") - return propagation_ids - - # Get recent agent memories from this task - recent_memories = await conn.fetch( - """ - SELECT * FROM agent_memory - WHERE agent_id = $1 - AND task_id = $2 - AND confidence_score >= 0.6 - AND propagated_to_team = FALSE - ORDER BY confidence_score DESC, created_at DESC - """, - agent_id, - task_id, - ) - - # Group memories by type and analyze for propagation - memory_groups = self._group_memories_for_propagation(recent_memories) - - for group_type, memories in memory_groups.items(): - if len(memories) >= 1 and self._meets_propagation_criteria( - memories, task_outcome - ): - # Create propagation task - propagation_id = await self._create_propagation_task( - source_type="agent", - source_id=agent_id, - target_type="team", - target_id=str(team_id), - knowledge_type=group_type, - knowledge_content_ids=[str(mem["id"]) for mem in memories], - propagation_trigger=PropagationTrigger.TASK_COMPLETION, - confidence_score=self._calculate_group_confidence(memories), - metadata={ - "task_id": task_id, - "task_outcome": task_outcome, - "memory_count": len(memories), - }, - ) - - propagation_ids.append(propagation_id) - - logger.info( - f"Created {len(propagation_ids)} propagation tasks for agent {agent_id} → team {team_id}" - ) - return propagation_ids - - async def trigger_team_to_org_propagation( - self, team_id: str, knowledge_threshold_check: bool = True - ) -> List[str]: - """Trigger knowledge propagation from team to organization level""" - - propagation_ids = [] - - async with self.pool.acquire() as conn: - # Get organization ID - org_id = await conn.fetchval( - """ - SELECT organization_id FROM teams WHERE id = $1 - """, - team_id, - ) - - if not org_id: - logger.warning(f"No organization found for team {team_id}") - return propagation_ids - - # Find high-value team knowledge for propagation - if knowledge_threshold_check: - team_knowledge = await conn.fetch( - """ - SELECT * FROM team_knowledge_base - WHERE team_id = $1 - AND effectiveness_score >= 0.7 - AND agent_adoption_rate >= 0.5 - AND created_at <= NOW() - INTERVAL '7 days' -- Allow time for validation - ORDER BY effectiveness_score DESC, agent_adoption_rate DESC - """, - team_id, - ) - else: - team_knowledge = await conn.fetch( - """ - SELECT * FROM team_knowledge_base - WHERE team_id = $1 - ORDER BY effectiveness_score DESC - LIMIT 10 - """, - team_id, - ) - - for knowledge in team_knowledge: - # Check if similar knowledge already exists at org level - if not await self._check_for_similar_org_knowledge( - knowledge, str(org_id) - ): - # Create propagation task - propagation_id = await self._create_propagation_task( - source_type="team", - source_id=team_id, - target_type="organization", - target_id=str(org_id), - knowledge_type=knowledge["knowledge_category"], - knowledge_content_ids=[str(knowledge["id"])], - propagation_trigger=PropagationTrigger.KNOWLEDGE_THRESHOLD, - confidence_score=knowledge["effectiveness_score"], - metadata={ - "team_knowledge_id": str(knowledge["id"]), - "adoption_rate": knowledge["agent_adoption_rate"], - "contributing_agents": knowledge["contributing_agents"], - }, - ) - - propagation_ids.append(propagation_id) - - logger.info( - f"Created {len(propagation_ids)} propagation tasks for team {team_id} → organization {org_id}" - ) - return propagation_ids - - async def trigger_cross_team_sharing( - self, - source_team_id: str, - knowledge_categories: List[KnowledgeCategory], - target_teams: Optional[List[str]] = None, - ) -> List[str]: - """Trigger knowledge sharing between teams""" - - propagation_ids = [] - - async with self.pool.acquire() as conn: - # Get organization and determine target teams - org_id = await conn.fetchval( - """ - SELECT organization_id FROM teams WHERE id = $1 - """, - source_team_id, - ) - - if not org_id: - return propagation_ids - - if not target_teams: - # Get all teams in the organization except source team - target_teams_rows = await conn.fetch( - """ - SELECT id FROM teams - WHERE organization_id = $1 AND id != $2 - """, - org_id, - source_team_id, - ) - target_teams = [str(row["id"]) for row in target_teams_rows] - - # Get relevant knowledge from source team - category_list = [cat.value for cat in knowledge_categories] - source_knowledge = await conn.fetch( - """ - SELECT * FROM team_knowledge_base - WHERE team_id = $1 - AND knowledge_category = ANY($2) - AND effectiveness_score >= 0.6 - ORDER BY effectiveness_score DESC - LIMIT 20 - """, - source_team_id, - category_list, - ) - - # Create propagation tasks for each target team - for target_team_id in target_teams: - for knowledge in source_knowledge: - # Check if target team would benefit from this knowledge - relevance = await self._calculate_cross_team_relevance( - knowledge, source_team_id, target_team_id - ) - - if relevance >= 0.5: - propagation_id = await self._create_propagation_task( - source_type="team", - source_id=source_team_id, - target_type="team", - target_id=target_team_id, - knowledge_type=knowledge["knowledge_category"], - knowledge_content_ids=[str(knowledge["id"])], - propagation_trigger=PropagationTrigger.CROSS_TEAM_REQUEST, - confidence_score=relevance, - metadata={ - "cross_team_relevance": relevance, - "source_effectiveness": knowledge[ - "effectiveness_score" - ], - }, - ) - - propagation_ids.append(propagation_id) - - logger.info(f"Created {len(propagation_ids)} cross-team propagation tasks") - return propagation_ids - - async def process_pending_propagations(self, limit: int = 10) -> Dict[str, int]: - """Process pending propagation tasks""" - - results = {"processed": 0, "completed": 0, "failed": 0} - - async with self.pool.acquire() as conn: - # Get pending propagation tasks - pending_tasks = await conn.fetch( - """ - SELECT * FROM knowledge_propagation_log - WHERE propagation_status = 'pending' - ORDER BY propagated_at ASC - LIMIT $1 - """, - limit, - ) - - for task_row in pending_tasks: - task = self._row_to_propagation_task(task_row) - - try: - # Update status to processing - await conn.execute( - """ - UPDATE knowledge_propagation_log - SET propagation_status = 'processing', processed_at = NOW() - WHERE id = $1 - """, - task.id, - ) - - # Process the propagation - success = await self._execute_propagation(task) - - if success: - # Mark as completed - await conn.execute( - """ - UPDATE knowledge_propagation_log - SET propagation_status = 'completed', - acceptance_status = 'accepted', - completed_at = NOW() - WHERE id = $1 - """, - task.id, - ) - results["completed"] += 1 - self.propagations_completed += 1 - else: - # Mark as failed - await conn.execute( - """ - UPDATE knowledge_propagation_log - SET propagation_status = 'failed' - WHERE id = $1 - """, - task.id, - ) - results["failed"] += 1 - - results["processed"] += 1 - self.propagations_processed += 1 - - except Exception as e: - logger.error(f"Error processing propagation task {task.id}: {e}") - await conn.execute( - """ - UPDATE knowledge_propagation_log - SET propagation_status = 'failed', - metadata = metadata || $2 - WHERE id = $1 - """, - task.id, - json.dumps({"error": str(e)}), - ) - results["failed"] += 1 - - return results - - async def get_propagation_statistics( - self, - organization_id: Optional[str] = None, - team_id: Optional[str] = None, - days_back: int = 30, - ) -> Dict[str, Any]: - """Get comprehensive propagation statistics""" - - async with self.pool.acquire() as conn: - where_conditions = [] - params = [] - - if organization_id: - where_conditions.append( - "target_id = $1 AND target_type = 'organization'" - ) - params.append(organization_id) - elif team_id: - where_conditions.append( - "(target_id = $1 OR source_id = $1) AND ('team' = ANY(ARRAY[target_type, source_type]))" - ) - params.append(team_id) - - # Parameterize the time window; days_back is bound, not interpolated. - days_param_idx = len(params) + 1 - params.append(days_back) - where_conditions.append( - f"propagated_at >= NOW() - (INTERVAL '1 day' * ${days_param_idx})" - ) - - where_clause = "WHERE " + " AND ".join(where_conditions) - - # Basic statistics - stats = await conn.fetchrow( - f""" - SELECT - COUNT(*) as total_propagations, - COUNT(CASE WHEN propagation_status = 'completed' THEN 1 END) as completed, - COUNT(CASE WHEN propagation_status = 'failed' THEN 1 END) as failed, - COUNT(CASE WHEN propagation_status = 'pending' THEN 1 END) as pending, - COUNT(CASE WHEN acceptance_status = 'accepted' THEN 1 END) as accepted, - COUNT(CASE WHEN acceptance_status = 'rejected' THEN 1 END) as rejected, - AVG(confidence_score) as avg_confidence - FROM knowledge_propagation_log - {where_clause} - """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params - *params, - ) - - # Propagation flow statistics - flow_stats = await conn.fetch( - f""" - SELECT - source_type || ' → ' || target_type as flow_type, - COUNT(*) as count, - AVG(confidence_score) as avg_confidence, - COUNT(CASE WHEN propagation_status = 'completed' THEN 1 END)::float / COUNT(*) as success_rate - FROM knowledge_propagation_log - {where_clause} - GROUP BY source_type, target_type - ORDER BY count DESC - """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params - *params, - ) - - # Trigger analysis - trigger_stats = await conn.fetch( - f""" - SELECT - propagation_trigger, - COUNT(*) as count, - AVG(confidence_score) as avg_confidence - FROM knowledge_propagation_log - {where_clause} - GROUP BY propagation_trigger - ORDER BY count DESC - """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params - *params, - ) - - return { - "time_period_days": days_back, - "basic_stats": dict(stats) if stats else {}, - "flow_patterns": [dict(flow) for flow in flow_stats], - "trigger_analysis": [dict(trigger) for trigger in trigger_stats], - "generated_at": datetime.now().isoformat(), - } - - async def _background_propagation_processor(self): - """Background task to continuously process propagation queue""" - - while self._running: - try: - # Process a batch of propagations - results = await self.process_pending_propagations( - self.propagation_batch_size - ) - - if results["processed"] > 0: - logger.info( - f"Processed {results['processed']} propagations: " - f"{results['completed']} completed, {results['failed']} failed" - ) - - # Sleep between processing cycles - await asyncio.sleep(30) # Process every 30 seconds - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Error in background propagation processor: {e}") - await asyncio.sleep(60) # Wait longer after errors - - async def _create_propagation_task( - self, - source_type: str, - source_id: str, - target_type: str, - target_id: str, - knowledge_type: str, - knowledge_content_ids: List[str], - propagation_trigger: PropagationTrigger, - confidence_score: float, - metadata: Dict[str, Any], - ) -> str: - """Create a new propagation task""" - - propagation_id = str(uuid.uuid4()) - - async with self.pool.acquire() as conn: - await conn.execute( - """ - INSERT INTO knowledge_propagation_log ( - id, source_type, source_id, target_type, target_id, - knowledge_type, propagation_method, propagation_trigger, - confidence_score, propagation_status, acceptance_status, metadata - ) VALUES ($1, $2, $3, $4, $5, $6, 'automatic', $7, $8, 'pending', 'pending', $9) - """, - propagation_id, - source_type, - source_id, - target_type, - target_id, - knowledge_type, - propagation_trigger.value, - confidence_score, - json.dumps({**metadata, "content_ids": knowledge_content_ids}), - ) - - return propagation_id - - async def _execute_propagation(self, task: PropagationTask) -> bool: - """Execute a specific propagation task""" - - try: - if task.source_type == "agent" and task.target_type == "team": - return await self._execute_agent_to_team_propagation(task) - elif task.source_type == "team" and task.target_type == "organization": - return await self._execute_team_to_org_propagation(task) - elif task.source_type == "team" and task.target_type == "team": - return await self._execute_team_to_team_propagation(task) - else: - logger.warning( - f"Unsupported propagation type: {task.source_type} → {task.target_type}" - ) - return False - - except Exception as e: - logger.error(f"Error executing propagation {task.id}: {e}") - return False - - async def _execute_agent_to_team_propagation(self, task: PropagationTask) -> bool: - """Execute agent → team propagation""" - - content_ids = task.metadata.get("content_ids", []) - if not content_ids: - return False - - # Aggregate agent memories to team knowledge - result = await self.team_knowledge_manager.aggregate_agent_knowledge_to_team( - team_id=task.target_id, - agent_id=task.source_id, - agent_memory_ids=content_ids, - aggregation_method="propagation", - ) - - return result is not None - - async def _execute_team_to_org_propagation(self, task: PropagationTask) -> bool: - """Execute team → organization propagation""" - - team_knowledge_id = task.metadata.get("team_knowledge_id") - if not team_knowledge_id: - return False - - async with self.pool.acquire() as conn: - # Get team knowledge - team_knowledge = await conn.fetchrow( - """ - SELECT * FROM team_knowledge_base WHERE id = $1 - """, - team_knowledge_id, - ) - - if not team_knowledge: - return False - - # Create organization knowledge - org_knowledge_id = await self.org_rag_manager.add_knowledge( - organization_id=task.target_id, - title=f"[Team Contribution] {team_knowledge['title']}", - content=team_knowledge["content"], - content_type=ContentType(team_knowledge["content_type"]), - knowledge_category=KnowledgeCategory( - team_knowledge["knowledge_category"] - ), - source_type=SourceType.TEAM_AGGREGATION, - source_team_id=task.source_id, - relevance_score=team_knowledge["effectiveness_score"], - quality_score=team_knowledge["effectiveness_score"], - visibility_level=VisibilityLevel.ORGANIZATION, - metadata={ - "team_adoption_rate": team_knowledge["agent_adoption_rate"], - "team_effectiveness": team_knowledge["effectiveness_score"], - "contributing_agents": team_knowledge["contributing_agents"], - "propagation_task_id": task.id, - }, - tags=team_knowledge["tags"] + ["team_contribution"], - ) - - return org_knowledge_id is not None - - async def _execute_team_to_team_propagation(self, task: PropagationTask) -> bool: - """Execute team → team propagation""" - - team_knowledge_id = task.metadata.get("content_ids", [None])[0] - if not team_knowledge_id: - return False - - async with self.pool.acquire() as conn: - # Get source team knowledge - source_knowledge = await conn.fetchrow( - """ - SELECT * FROM team_knowledge_base WHERE id = $1 - """, - team_knowledge_id, - ) - - if not source_knowledge: - return False - - # Create adapted knowledge for target team - adapted_knowledge_id = ( - await self.team_knowledge_manager.create_team_knowledge( - team_id=task.target_id, - title=f"[Shared] {source_knowledge['title']}", - content=source_knowledge["content"], - content_type=ContentType(source_knowledge["content_type"]), - knowledge_category=KnowledgeCategory( - source_knowledge["knowledge_category"] - ), - source_type=SourceType.TEAM_AGGREGATION, - contributing_agents=[], - source_knowledge_ids=[team_knowledge_id], - aggregation_method="cross_team_sharing", - team_relevance_score=task.confidence_score, - metadata={ - "source_team_id": task.source_id, - "cross_team_propagation": True, - "original_effectiveness": source_knowledge[ - "effectiveness_score" - ], - "propagation_task_id": task.id, - }, - tags=source_knowledge["tags"] + ["cross_team_shared"], - ) - ) - - return adapted_knowledge_id is not None - - def _group_memories_for_propagation(self, memories: List) -> Dict[str, List]: - """Group memories by type for propagation analysis""" - groups = {} - - for memory in memories: - memory_type = memory.get("memory_type", "general") - if memory_type not in groups: - groups[memory_type] = [] - groups[memory_type].append(memory) - - return groups - - def _meets_propagation_criteria( - self, memories: List, task_outcome: Dict[str, Any] - ) -> bool: - """Determine if memories meet criteria for propagation""" - - # Check success rate - success = task_outcome.get("success", False) - if not success: - return False - - # Check confidence scores - avg_confidence = sum(mem["confidence_score"] for mem in memories) / len( - memories - ) - if avg_confidence < 0.6: - return False - - # Check memory age (don't propagate very old memories) - recent_memories = [ - mem for mem in memories if (datetime.now() - mem["created_at"]).days <= 7 - ] - - return ( - len(recent_memories) >= len(memories) * 0.5 - ) # At least 50% should be recent - - def _calculate_group_confidence(self, memories: List) -> float: - """Calculate confidence score for a group of memories""" - if not memories: - return 0.0 - - scores = [mem["confidence_score"] for mem in memories] - success_correlations = [mem.get("success_correlation", 0.0) for mem in memories] - - # Weighted average with recency bias - weights = [1.0 / (1 + i * 0.1) for i in range(len(memories))] - - weighted_confidence = sum(s * w for s, w in zip(scores, weights)) / sum(weights) - avg_success = ( - sum(success_correlations) / len(success_correlations) - if success_correlations - else 0.0 - ) - - return min(1.0, weighted_confidence * 0.7 + avg_success * 0.3) - - async def _check_for_similar_org_knowledge( - self, team_knowledge: Dict, org_id: str - ) -> bool: - """Check if similar knowledge already exists at organization level""" - - if not team_knowledge.get("embedding"): - return False - - # Search for similar content - search_results = await self.org_rag_manager.search_knowledge( - organization_id=org_id, - query=team_knowledge["content"][:200], # Use beginning of content as query - limit=5, - min_similarity=self.similarity_threshold, - ) - - # Check if any results are highly similar - for result in search_results: - if result.similarity_score >= self.similarity_threshold: - return True - - return False - - async def _calculate_cross_team_relevance( - self, knowledge: Dict, source_team_id: str, target_team_id: str - ) -> float: - """Calculate how relevant knowledge from one team is for another team""" - - async with self.pool.acquire() as conn: - # Get team information - teams = await conn.fetch( - """ - SELECT id, team_type, settings FROM teams - WHERE id IN ($1, $2) - """, - source_team_id, - target_team_id, - ) - - if len(teams) != 2: - return 0.0 - - source_team = next(t for t in teams if str(t["id"]) == source_team_id) - target_team = next(t for t in teams if str(t["id"]) == target_team_id) - - relevance_factors = [] - - # Factor 1: Team type similarity - type_similarity = ( - 1.0 if source_team["team_type"] == target_team["team_type"] else 0.3 - ) - relevance_factors.append(type_similarity) - - # Factor 2: Knowledge category relevance to target team - target_team_categories = await conn.fetch( - """ - SELECT knowledge_category, COUNT(*) as usage_count - FROM team_knowledge_base - WHERE team_id = $1 - GROUP BY knowledge_category - ORDER BY usage_count DESC - LIMIT 5 - """, - target_team_id, - ) - - target_categories = [ - cat["knowledge_category"] for cat in target_team_categories - ] - category_relevance = ( - 1.0 if knowledge["knowledge_category"] in target_categories else 0.4 - ) - relevance_factors.append(category_relevance) - - # Factor 3: Effectiveness score of source knowledge - effectiveness_factor = knowledge["effectiveness_score"] - relevance_factors.append(effectiveness_factor) - - # Factor 4: Adoption rate in source team (indicates broad utility) - adoption_factor = knowledge["agent_adoption_rate"] - relevance_factors.append(adoption_factor) - - # Calculate weighted relevance - weights = [0.2, 0.3, 0.3, 0.2] - relevance = sum(f * w for f, w in zip(relevance_factors, weights)) - - return min(1.0, max(0.0, relevance)) - - def _create_default_propagation_rules(self) -> List[PropagationRule]: - """Create default propagation rules""" - return [ - # Agent to Team rules - PropagationRule( - source_type="agent", - target_type="team", - min_confidence=0.6, - min_success_correlation=0.0, - min_usage_count=1, - knowledge_categories=list(KnowledgeCategory), - auto_approve=True, - propagation_weight=1.0, - ), - # Team to Organization rules - PropagationRule( - source_type="team", - target_type="organization", - min_confidence=0.7, - min_success_correlation=0.5, - min_usage_count=3, - knowledge_categories=list(KnowledgeCategory), - auto_approve=False, - propagation_weight=0.8, - ), - # Cross-team rules - PropagationRule( - source_type="team", - target_type="team", - min_confidence=0.65, - min_success_correlation=0.4, - min_usage_count=2, - knowledge_categories=[ - KnowledgeCategory.DEVELOPMENT, - KnowledgeCategory.TESTING, - KnowledgeCategory.TROUBLESHOOTING, - ], - auto_approve=False, - propagation_weight=0.6, - ), - ] - - def _row_to_propagation_task(self, row) -> PropagationTask: - """Convert database row to PropagationTask object""" - return PropagationTask( - id=str(row["id"]), - source_type=row["source_type"], - source_id=str(row["source_id"]), - target_type=row["target_type"], - target_id=str(row["target_id"]), - knowledge_type=row["knowledge_type"], - knowledge_content_id=( - str(row["knowledge_content_id"]) if row["knowledge_content_id"] else "" - ), - propagation_method=row["propagation_method"], - propagation_trigger=PropagationTrigger(row["propagation_trigger"]), - confidence_score=row["confidence_score"], - propagation_status=PropagationStatus(row["propagation_status"]), - acceptance_status=AcceptanceStatus(row["acceptance_status"]), - metadata=( - json.loads(row["metadata"]) - if isinstance(row["metadata"], str) - else row["metadata"] - ), - created_at=row["propagated_at"], - processed_at=row["processed_at"], - completed_at=row["completed_at"], - ) +""" +Knowledge Propagation Engine for FuzeAgent + +This module handles automated knowledge flow between agents, teams, and organizations. +It determines when knowledge should be propagated, executes the propagation, +and manages the lifecycle of knowledge across hierarchical levels. +""" + +import asyncio +import json +import logging +import uuid +from dataclasses import dataclass +from datetime import datetime, timedelta +from enum import Enum +from typing import Any, Dict, List, Optional, Set, Tuple + +import asyncpg +from sentence_transformers import SentenceTransformer + +from .organization_rag_manager import ( + ContentType, + KnowledgeCategory, + OrganizationRAGManager, + SourceType, + VisibilityLevel, +) +from .team_knowledge_manager import TeamKnowledgeManager + +logger = logging.getLogger(__name__) + + +class PropagationTrigger(str, Enum): + TASK_COMPLETION = "task_completion" + KNOWLEDGE_THRESHOLD = "knowledge_threshold" + MANUAL_REQUEST = "manual_request" + SCHEDULED_SYNC = "scheduled_sync" + CROSS_TEAM_REQUEST = "cross_team_request" + QUALITY_IMPROVEMENT = "quality_improvement" + + +class PropagationStatus(str, Enum): + PENDING = "pending" + PROCESSING = "processing" + COMPLETED = "completed" + FAILED = "failed" + REJECTED = "rejected" + + +class AcceptanceStatus(str, Enum): + PENDING = "pending" + ACCEPTED = "accepted" + REJECTED = "rejected" + MODIFIED = "modified" + + +@dataclass +class PropagationRule: + """Defines rules for knowledge propagation""" + + source_type: str # 'agent', 'team', 'organization' + target_type: str # 'agent', 'team', 'organization' + min_confidence: float + min_success_correlation: float + min_usage_count: int + knowledge_categories: List[KnowledgeCategory] + auto_approve: bool + propagation_weight: float + + +@dataclass +class PropagationTask: + """Represents a knowledge propagation task""" + + id: str + source_type: str + source_id: str + target_type: str + target_id: str + knowledge_type: str + knowledge_content_id: str + propagation_method: str + propagation_trigger: PropagationTrigger + confidence_score: float + propagation_status: PropagationStatus + acceptance_status: AcceptanceStatus + metadata: Dict[str, Any] + created_at: datetime + processed_at: Optional[datetime] + completed_at: Optional[datetime] + + +class KnowledgePropagationEngine: + """ + Manages the automated flow of knowledge across the organization hierarchy. + Handles agent → team → organization propagation and cross-team sharing. + """ + + def __init__( + self, + database_url: str, + org_rag_manager: OrganizationRAGManager, + team_knowledge_manager: TeamKnowledgeManager, + ): + self.database_url = database_url + self.org_rag_manager = org_rag_manager + self.team_knowledge_manager = team_knowledge_manager + self.pool: Optional[asyncpg.Pool] = None + + # Initialize embedding model for similarity analysis + self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") + + # Default propagation rules + self.default_rules = self._create_default_propagation_rules() + + # Configuration + self.propagation_batch_size = 50 + self.max_concurrent_propagations = 5 + self.similarity_threshold = 0.8 + self.propagation_cooldown_hours = 24 + + # Statistics + self.propagations_processed = 0 + self.propagations_completed = 0 + self.propagations_rejected = 0 + + # Background task management + self._propagation_task: Optional[asyncio.Task] = None + self._running = False + + async def initialize(self): + """Initialize the knowledge propagation engine""" + logger.info("Initializing KnowledgePropagationEngine") + + try: + self.pool = await asyncpg.create_pool( + self.database_url, min_size=2, max_size=10, command_timeout=60 + ) + + # Start background propagation processing + self._running = True + self._propagation_task = asyncio.create_task( + self._background_propagation_processor() + ) + + logger.info("KnowledgePropagationEngine initialized successfully") + + except Exception as e: + logger.error(f"Failed to initialize KnowledgePropagationEngine: {e}") + raise + + async def close(self): + """Close the propagation engine and cleanup resources""" + self._running = False + + if self._propagation_task: + self._propagation_task.cancel() + try: + await self._propagation_task + except asyncio.CancelledError: + pass + + if self.pool: + await self.pool.close() + + logger.info("KnowledgePropagationEngine closed") + + async def trigger_agent_to_team_propagation( + self, agent_id: str, task_id: str, task_outcome: Dict[str, Any] + ) -> List[str]: + """Trigger knowledge propagation from agent to team level after task completion""" + + propagation_ids = [] + + async with self.pool.acquire() as conn: + # Get agent's team + team_id = await conn.fetchval( + """ + SELECT team_id FROM agents WHERE id = $1 + """, + agent_id, + ) + + if not team_id: + logger.warning(f"No team found for agent {agent_id}") + return propagation_ids + + # Get recent agent memories from this task + recent_memories = await conn.fetch( + """ + SELECT * FROM agent_memory + WHERE agent_id = $1 + AND task_id = $2 + AND confidence_score >= 0.6 + AND propagated_to_team = FALSE + ORDER BY confidence_score DESC, created_at DESC + """, + agent_id, + task_id, + ) + + # Group memories by type and analyze for propagation + memory_groups = self._group_memories_for_propagation(recent_memories) + + for group_type, memories in memory_groups.items(): + if len(memories) >= 1 and self._meets_propagation_criteria( + memories, task_outcome + ): + # Create propagation task + propagation_id = await self._create_propagation_task( + source_type="agent", + source_id=agent_id, + target_type="team", + target_id=str(team_id), + knowledge_type=group_type, + knowledge_content_ids=[str(mem["id"]) for mem in memories], + propagation_trigger=PropagationTrigger.TASK_COMPLETION, + confidence_score=self._calculate_group_confidence(memories), + metadata={ + "task_id": task_id, + "task_outcome": task_outcome, + "memory_count": len(memories), + }, + ) + + propagation_ids.append(propagation_id) + + logger.info( + f"Created {len(propagation_ids)} propagation tasks for agent {agent_id} → team {team_id}" + ) + return propagation_ids + + async def trigger_team_to_org_propagation( + self, team_id: str, knowledge_threshold_check: bool = True + ) -> List[str]: + """Trigger knowledge propagation from team to organization level""" + + propagation_ids = [] + + async with self.pool.acquire() as conn: + # Get organization ID + org_id = await conn.fetchval( + """ + SELECT organization_id FROM teams WHERE id = $1 + """, + team_id, + ) + + if not org_id: + logger.warning(f"No organization found for team {team_id}") + return propagation_ids + + # Find high-value team knowledge for propagation + if knowledge_threshold_check: + team_knowledge = await conn.fetch( + """ + SELECT * FROM team_knowledge_base + WHERE team_id = $1 + AND effectiveness_score >= 0.7 + AND agent_adoption_rate >= 0.5 + AND created_at <= NOW() - INTERVAL '7 days' -- Allow time for validation + ORDER BY effectiveness_score DESC, agent_adoption_rate DESC + """, + team_id, + ) + else: + team_knowledge = await conn.fetch( + """ + SELECT * FROM team_knowledge_base + WHERE team_id = $1 + ORDER BY effectiveness_score DESC + LIMIT 10 + """, + team_id, + ) + + for knowledge in team_knowledge: + # Check if similar knowledge already exists at org level + if not await self._check_for_similar_org_knowledge( + knowledge, str(org_id) + ): + # Create propagation task + propagation_id = await self._create_propagation_task( + source_type="team", + source_id=team_id, + target_type="organization", + target_id=str(org_id), + knowledge_type=knowledge["knowledge_category"], + knowledge_content_ids=[str(knowledge["id"])], + propagation_trigger=PropagationTrigger.KNOWLEDGE_THRESHOLD, + confidence_score=knowledge["effectiveness_score"], + metadata={ + "team_knowledge_id": str(knowledge["id"]), + "adoption_rate": knowledge["agent_adoption_rate"], + "contributing_agents": knowledge["contributing_agents"], + }, + ) + + propagation_ids.append(propagation_id) + + logger.info( + f"Created {len(propagation_ids)} propagation tasks for team {team_id} → organization {org_id}" + ) + return propagation_ids + + async def trigger_cross_team_sharing( + self, + source_team_id: str, + knowledge_categories: List[KnowledgeCategory], + target_teams: Optional[List[str]] = None, + ) -> List[str]: + """Trigger knowledge sharing between teams""" + + propagation_ids = [] + + async with self.pool.acquire() as conn: + # Get organization and determine target teams + org_id = await conn.fetchval( + """ + SELECT organization_id FROM teams WHERE id = $1 + """, + source_team_id, + ) + + if not org_id: + return propagation_ids + + if not target_teams: + # Get all teams in the organization except source team + target_teams_rows = await conn.fetch( + """ + SELECT id FROM teams + WHERE organization_id = $1 AND id != $2 + """, + org_id, + source_team_id, + ) + target_teams = [str(row["id"]) for row in target_teams_rows] + + # Get relevant knowledge from source team + category_list = [cat.value for cat in knowledge_categories] + source_knowledge = await conn.fetch( + """ + SELECT * FROM team_knowledge_base + WHERE team_id = $1 + AND knowledge_category = ANY($2) + AND effectiveness_score >= 0.6 + ORDER BY effectiveness_score DESC + LIMIT 20 + """, + source_team_id, + category_list, + ) + + # Create propagation tasks for each target team + for target_team_id in target_teams: + for knowledge in source_knowledge: + # Check if target team would benefit from this knowledge + relevance = await self._calculate_cross_team_relevance( + knowledge, source_team_id, target_team_id + ) + + if relevance >= 0.5: + propagation_id = await self._create_propagation_task( + source_type="team", + source_id=source_team_id, + target_type="team", + target_id=target_team_id, + knowledge_type=knowledge["knowledge_category"], + knowledge_content_ids=[str(knowledge["id"])], + propagation_trigger=PropagationTrigger.CROSS_TEAM_REQUEST, + confidence_score=relevance, + metadata={ + "cross_team_relevance": relevance, + "source_effectiveness": knowledge[ + "effectiveness_score" + ], + }, + ) + + propagation_ids.append(propagation_id) + + logger.info(f"Created {len(propagation_ids)} cross-team propagation tasks") + return propagation_ids + + async def process_pending_propagations(self, limit: int = 10) -> Dict[str, int]: + """Process pending propagation tasks""" + + results = {"processed": 0, "completed": 0, "failed": 0} + + async with self.pool.acquire() as conn: + # Get pending propagation tasks + pending_tasks = await conn.fetch( + """ + SELECT * FROM knowledge_propagation_log + WHERE propagation_status = 'pending' + ORDER BY propagated_at ASC + LIMIT $1 + """, + limit, + ) + + for task_row in pending_tasks: + task = self._row_to_propagation_task(task_row) + + try: + # Update status to processing + await conn.execute( + """ + UPDATE knowledge_propagation_log + SET propagation_status = 'processing', processed_at = NOW() + WHERE id = $1 + """, + task.id, + ) + + # Process the propagation + success = await self._execute_propagation(task) + + if success: + # Mark as completed + await conn.execute( + """ + UPDATE knowledge_propagation_log + SET propagation_status = 'completed', + acceptance_status = 'accepted', + completed_at = NOW() + WHERE id = $1 + """, + task.id, + ) + results["completed"] += 1 + self.propagations_completed += 1 + else: + # Mark as failed + await conn.execute( + """ + UPDATE knowledge_propagation_log + SET propagation_status = 'failed' + WHERE id = $1 + """, + task.id, + ) + results["failed"] += 1 + + results["processed"] += 1 + self.propagations_processed += 1 + + except Exception as e: + logger.error(f"Error processing propagation task {task.id}: {e}") + await conn.execute( + """ + UPDATE knowledge_propagation_log + SET propagation_status = 'failed', + metadata = metadata || $2 + WHERE id = $1 + """, + task.id, + json.dumps({"error": str(e)}), + ) + results["failed"] += 1 + + return results + + async def get_propagation_statistics( + self, + organization_id: Optional[str] = None, + team_id: Optional[str] = None, + days_back: int = 30, + ) -> Dict[str, Any]: + """Get comprehensive propagation statistics""" + + async with self.pool.acquire() as conn: + where_conditions = [] + params = [] + + if organization_id: + where_conditions.append( + "target_id = $1 AND target_type = 'organization'" + ) + params.append(organization_id) + elif team_id: + where_conditions.append( + "(target_id = $1 OR source_id = $1) AND ('team' = ANY(ARRAY[target_type, source_type]))" + ) + params.append(team_id) + + # Parameterize the time window; days_back is bound, not interpolated. + days_param_idx = len(params) + 1 + params.append(days_back) + where_conditions.append( + f"propagated_at >= NOW() - (INTERVAL '1 day' * ${days_param_idx})" + ) + + where_clause = "WHERE " + " AND ".join(where_conditions) + + # Basic statistics + stats = await conn.fetchrow( + f""" + SELECT + COUNT(*) as total_propagations, + COUNT(CASE WHEN propagation_status = 'completed' THEN 1 END) as completed, + COUNT(CASE WHEN propagation_status = 'failed' THEN 1 END) as failed, + COUNT(CASE WHEN propagation_status = 'pending' THEN 1 END) as pending, + COUNT(CASE WHEN acceptance_status = 'accepted' THEN 1 END) as accepted, + COUNT(CASE WHEN acceptance_status = 'rejected' THEN 1 END) as rejected, + AVG(confidence_score) as avg_confidence + FROM knowledge_propagation_log + {where_clause} + """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params + *params, + ) + + # Propagation flow statistics + flow_stats = await conn.fetch( + f""" + SELECT + source_type || ' → ' || target_type as flow_type, + COUNT(*) as count, + AVG(confidence_score) as avg_confidence, + COUNT(CASE WHEN propagation_status = 'completed' THEN 1 END)::float / COUNT(*) as success_rate + FROM knowledge_propagation_log + {where_clause} + GROUP BY source_type, target_type + ORDER BY count DESC + """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params + *params, + ) + + # Trigger analysis + trigger_stats = await conn.fetch( + f""" + SELECT + propagation_trigger, + COUNT(*) as count, + AVG(confidence_score) as avg_confidence + FROM knowledge_propagation_log + {where_clause} + GROUP BY propagation_trigger + ORDER BY count DESC + """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params + *params, + ) + + return { + "time_period_days": days_back, + "basic_stats": dict(stats) if stats else {}, + "flow_patterns": [dict(flow) for flow in flow_stats], + "trigger_analysis": [dict(trigger) for trigger in trigger_stats], + "generated_at": datetime.now().isoformat(), + } + + async def _background_propagation_processor(self): + """Background task to continuously process propagation queue""" + + while self._running: + try: + # Process a batch of propagations + results = await self.process_pending_propagations( + self.propagation_batch_size + ) + + if results["processed"] > 0: + logger.info( + f"Processed {results['processed']} propagations: " + f"{results['completed']} completed, {results['failed']} failed" + ) + + # Sleep between processing cycles + await asyncio.sleep(30) # Process every 30 seconds + + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"Error in background propagation processor: {e}") + await asyncio.sleep(60) # Wait longer after errors + + async def _create_propagation_task( + self, + source_type: str, + source_id: str, + target_type: str, + target_id: str, + knowledge_type: str, + knowledge_content_ids: List[str], + propagation_trigger: PropagationTrigger, + confidence_score: float, + metadata: Dict[str, Any], + ) -> str: + """Create a new propagation task""" + + propagation_id = str(uuid.uuid4()) + + async with self.pool.acquire() as conn: + await conn.execute( + """ + INSERT INTO knowledge_propagation_log ( + id, source_type, source_id, target_type, target_id, + knowledge_type, propagation_method, propagation_trigger, + confidence_score, propagation_status, acceptance_status, metadata + ) VALUES ($1, $2, $3, $4, $5, $6, 'automatic', $7, $8, 'pending', 'pending', $9) + """, + propagation_id, + source_type, + source_id, + target_type, + target_id, + knowledge_type, + propagation_trigger.value, + confidence_score, + json.dumps({**metadata, "content_ids": knowledge_content_ids}), + ) + + return propagation_id + + async def _execute_propagation(self, task: PropagationTask) -> bool: + """Execute a specific propagation task""" + + try: + if task.source_type == "agent" and task.target_type == "team": + return await self._execute_agent_to_team_propagation(task) + elif task.source_type == "team" and task.target_type == "organization": + return await self._execute_team_to_org_propagation(task) + elif task.source_type == "team" and task.target_type == "team": + return await self._execute_team_to_team_propagation(task) + else: + logger.warning( + f"Unsupported propagation type: {task.source_type} → {task.target_type}" + ) + return False + + except Exception as e: + logger.error(f"Error executing propagation {task.id}: {e}") + return False + + async def _execute_agent_to_team_propagation(self, task: PropagationTask) -> bool: + """Execute agent → team propagation""" + + content_ids = task.metadata.get("content_ids", []) + if not content_ids: + return False + + # Aggregate agent memories to team knowledge + result = await self.team_knowledge_manager.aggregate_agent_knowledge_to_team( + team_id=task.target_id, + agent_id=task.source_id, + agent_memory_ids=content_ids, + aggregation_method="propagation", + ) + + return result is not None + + async def _execute_team_to_org_propagation(self, task: PropagationTask) -> bool: + """Execute team → organization propagation""" + + team_knowledge_id = task.metadata.get("team_knowledge_id") + if not team_knowledge_id: + return False + + async with self.pool.acquire() as conn: + # Get team knowledge + team_knowledge = await conn.fetchrow( + """ + SELECT * FROM team_knowledge_base WHERE id = $1 + """, + team_knowledge_id, + ) + + if not team_knowledge: + return False + + # Create organization knowledge + org_knowledge_id = await self.org_rag_manager.add_knowledge( + organization_id=task.target_id, + title=f"[Team Contribution] {team_knowledge['title']}", + content=team_knowledge["content"], + content_type=ContentType(team_knowledge["content_type"]), + knowledge_category=KnowledgeCategory( + team_knowledge["knowledge_category"] + ), + source_type=SourceType.TEAM_AGGREGATION, + source_team_id=task.source_id, + relevance_score=team_knowledge["effectiveness_score"], + quality_score=team_knowledge["effectiveness_score"], + visibility_level=VisibilityLevel.ORGANIZATION, + metadata={ + "team_adoption_rate": team_knowledge["agent_adoption_rate"], + "team_effectiveness": team_knowledge["effectiveness_score"], + "contributing_agents": team_knowledge["contributing_agents"], + "propagation_task_id": task.id, + }, + tags=team_knowledge["tags"] + ["team_contribution"], + ) + + return org_knowledge_id is not None + + async def _execute_team_to_team_propagation(self, task: PropagationTask) -> bool: + """Execute team → team propagation""" + + team_knowledge_id = task.metadata.get("content_ids", [None])[0] + if not team_knowledge_id: + return False + + async with self.pool.acquire() as conn: + # Get source team knowledge + source_knowledge = await conn.fetchrow( + """ + SELECT * FROM team_knowledge_base WHERE id = $1 + """, + team_knowledge_id, + ) + + if not source_knowledge: + return False + + # Create adapted knowledge for target team + adapted_knowledge_id = ( + await self.team_knowledge_manager.create_team_knowledge( + team_id=task.target_id, + title=f"[Shared] {source_knowledge['title']}", + content=source_knowledge["content"], + content_type=ContentType(source_knowledge["content_type"]), + knowledge_category=KnowledgeCategory( + source_knowledge["knowledge_category"] + ), + source_type=SourceType.TEAM_AGGREGATION, + contributing_agents=[], + source_knowledge_ids=[team_knowledge_id], + aggregation_method="cross_team_sharing", + team_relevance_score=task.confidence_score, + metadata={ + "source_team_id": task.source_id, + "cross_team_propagation": True, + "original_effectiveness": source_knowledge[ + "effectiveness_score" + ], + "propagation_task_id": task.id, + }, + tags=source_knowledge["tags"] + ["cross_team_shared"], + ) + ) + + return adapted_knowledge_id is not None + + def _group_memories_for_propagation(self, memories: List) -> Dict[str, List]: + """Group memories by type for propagation analysis""" + groups = {} + + for memory in memories: + memory_type = memory.get("memory_type", "general") + if memory_type not in groups: + groups[memory_type] = [] + groups[memory_type].append(memory) + + return groups + + def _meets_propagation_criteria( + self, memories: List, task_outcome: Dict[str, Any] + ) -> bool: + """Determine if memories meet criteria for propagation""" + + # Check success rate + success = task_outcome.get("success", False) + if not success: + return False + + # Check confidence scores + avg_confidence = sum(mem["confidence_score"] for mem in memories) / len( + memories + ) + if avg_confidence < 0.6: + return False + + # Check memory age (don't propagate very old memories) + recent_memories = [ + mem for mem in memories if (datetime.now() - mem["created_at"]).days <= 7 + ] + + return ( + len(recent_memories) >= len(memories) * 0.5 + ) # At least 50% should be recent + + def _calculate_group_confidence(self, memories: List) -> float: + """Calculate confidence score for a group of memories""" + if not memories: + return 0.0 + + scores = [mem["confidence_score"] for mem in memories] + success_correlations = [mem.get("success_correlation", 0.0) for mem in memories] + + # Weighted average with recency bias + weights = [1.0 / (1 + i * 0.1) for i in range(len(memories))] + + weighted_confidence = sum(s * w for s, w in zip(scores, weights)) / sum(weights) + avg_success = ( + sum(success_correlations) / len(success_correlations) + if success_correlations + else 0.0 + ) + + return min(1.0, weighted_confidence * 0.7 + avg_success * 0.3) + + async def _check_for_similar_org_knowledge( + self, team_knowledge: Dict, org_id: str + ) -> bool: + """Check if similar knowledge already exists at organization level""" + + if not team_knowledge.get("embedding"): + return False + + # Search for similar content + search_results = await self.org_rag_manager.search_knowledge( + organization_id=org_id, + query=team_knowledge["content"][:200], # Use beginning of content as query + limit=5, + min_similarity=self.similarity_threshold, + ) + + # Check if any results are highly similar + for result in search_results: + if result.similarity_score >= self.similarity_threshold: + return True + + return False + + async def _calculate_cross_team_relevance( + self, knowledge: Dict, source_team_id: str, target_team_id: str + ) -> float: + """Calculate how relevant knowledge from one team is for another team""" + + async with self.pool.acquire() as conn: + # Get team information + teams = await conn.fetch( + """ + SELECT id, team_type, settings FROM teams + WHERE id IN ($1, $2) + """, + source_team_id, + target_team_id, + ) + + if len(teams) != 2: + return 0.0 + + source_team = next(t for t in teams if str(t["id"]) == source_team_id) + target_team = next(t for t in teams if str(t["id"]) == target_team_id) + + relevance_factors = [] + + # Factor 1: Team type similarity + type_similarity = ( + 1.0 if source_team["team_type"] == target_team["team_type"] else 0.3 + ) + relevance_factors.append(type_similarity) + + # Factor 2: Knowledge category relevance to target team + target_team_categories = await conn.fetch( + """ + SELECT knowledge_category, COUNT(*) as usage_count + FROM team_knowledge_base + WHERE team_id = $1 + GROUP BY knowledge_category + ORDER BY usage_count DESC + LIMIT 5 + """, + target_team_id, + ) + + target_categories = [ + cat["knowledge_category"] for cat in target_team_categories + ] + category_relevance = ( + 1.0 if knowledge["knowledge_category"] in target_categories else 0.4 + ) + relevance_factors.append(category_relevance) + + # Factor 3: Effectiveness score of source knowledge + effectiveness_factor = knowledge["effectiveness_score"] + relevance_factors.append(effectiveness_factor) + + # Factor 4: Adoption rate in source team (indicates broad utility) + adoption_factor = knowledge["agent_adoption_rate"] + relevance_factors.append(adoption_factor) + + # Calculate weighted relevance + weights = [0.2, 0.3, 0.3, 0.2] + relevance = sum(f * w for f, w in zip(relevance_factors, weights)) + + return min(1.0, max(0.0, relevance)) + + def _create_default_propagation_rules(self) -> List[PropagationRule]: + """Create default propagation rules""" + return [ + # Agent to Team rules + PropagationRule( + source_type="agent", + target_type="team", + min_confidence=0.6, + min_success_correlation=0.0, + min_usage_count=1, + knowledge_categories=list(KnowledgeCategory), + auto_approve=True, + propagation_weight=1.0, + ), + # Team to Organization rules + PropagationRule( + source_type="team", + target_type="organization", + min_confidence=0.7, + min_success_correlation=0.5, + min_usage_count=3, + knowledge_categories=list(KnowledgeCategory), + auto_approve=False, + propagation_weight=0.8, + ), + # Cross-team rules + PropagationRule( + source_type="team", + target_type="team", + min_confidence=0.65, + min_success_correlation=0.4, + min_usage_count=2, + knowledge_categories=[ + KnowledgeCategory.DEVELOPMENT, + KnowledgeCategory.TESTING, + KnowledgeCategory.TROUBLESHOOTING, + ], + auto_approve=False, + propagation_weight=0.6, + ), + ] + + def _row_to_propagation_task(self, row) -> PropagationTask: + """Convert database row to PropagationTask object""" + return PropagationTask( + id=str(row["id"]), + source_type=row["source_type"], + source_id=str(row["source_id"]), + target_type=row["target_type"], + target_id=str(row["target_id"]), + knowledge_type=row["knowledge_type"], + knowledge_content_id=( + str(row["knowledge_content_id"]) if row["knowledge_content_id"] else "" + ), + propagation_method=row["propagation_method"], + propagation_trigger=PropagationTrigger(row["propagation_trigger"]), + confidence_score=row["confidence_score"], + propagation_status=PropagationStatus(row["propagation_status"]), + acceptance_status=AcceptanceStatus(row["acceptance_status"]), + metadata=( + json.loads(row["metadata"]) + if isinstance(row["metadata"], str) + else row["metadata"] + ), + created_at=row["propagated_at"], + processed_at=row["processed_at"], + completed_at=row["completed_at"], + ) diff --git a/services/orchestrator/main.py b/services/orchestrator/main.py index bb7fbe2..ce26700 100644 --- a/services/orchestrator/main.py +++ b/services/orchestrator/main.py @@ -1,6013 +1,6013 @@ -import asyncio -import json -import logging -import os -from collections import defaultdict -from contextlib import asynccontextmanager -from datetime import date, datetime -from decimal import Decimal -from typing import Any, Dict, List, Optional - -import jwt -from fastapi import ( - Body, - Depends, - FastAPI, - File, - Form, - HTTPException, - Path, - Query, - UploadFile, - WebSocket, - WebSocketDisconnect, - status, -) -from fastapi.middleware.cors import CORSMiddleware -from fastapi.responses import HTMLResponse, Response -from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer -from pydantic import BaseModel, Field - -from hierarchy_endpoints import router as hierarchy_router - -from .agent_manager import AgentManager -from .container_manager import ContainerConfig, ContainerStatus, container_manager -from .context_service import ContextService -from .database import get_db_connection -from .knowledge_manager import DocumentMetadata, knowledge_manager -from .rag_integration import RAGContext, rag_system -from .sandbox_manager import AgentSandboxManager -from .task_execution_engine import TaskExecutionEngine -from .task_queue import TaskQueue -from .websocket_manager import ( - UpdateType, - WebSocketUpdate, - notify_agent_status_change, - notify_container_status_change, - notify_knowledge_update, - notify_task_progress, - websocket_manager, -) - -logger = logging.getLogger(__name__) - -# --------------------------------------------------------------------------- -# Auth helpers (Track 3) -# --------------------------------------------------------------------------- -_security = HTTPBearer(auto_error=False) -_jwt_secret = os.environ.get("FUZEFRONT_JWT_SECRET", "") - - -def require_auth(credentials: HTTPAuthorizationCredentials = Depends(_security)): - """Verify FuzeFront JWT on mutating endpoints. Disabled when secret not set (dev).""" - if not _jwt_secret: - return None # Auth disabled when secret not configured (dev mode) - if not credentials: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing token" - ) - try: - payload = jwt.decode(credentials.credentials, _jwt_secret, algorithms=["HS256"]) - return payload - except jwt.ExpiredSignatureError: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, detail="Token expired" - ) - except jwt.InvalidTokenError: - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token" - ) - - -# --------------------------------------------------------------------------- -# Agent relay state (Track 4) -# --------------------------------------------------------------------------- -# agent_id -> list of subscriber WebSockets watching that agent's session -agent_relay_subscribers: Dict[str, List[WebSocket]] = defaultdict(list) - - -# Pydantic models for API documentation -class AgentCreateRequest(BaseModel): - name: str = Field(..., description="Agent name") - role: str = Field(..., description="Agent role (e.g., 'Senior React Developer')") - type: str = Field(..., description="Agent type (e.g., 'developer', 'executive')") - config: Dict[str, Any] = Field( - default_factory=dict, description="Agent configuration" - ) - repository_settings: Dict[str, Any] = Field( - default_factory=dict, description="Repository settings" - ) - sandbox_settings: Dict[str, Any] = Field( - default_factory=dict, description="Sandbox settings" - ) - - -class TaskCreateRequest(BaseModel): - title: str = Field(..., description="Task title") - description: str = Field(..., description="Task description") - priority: str = Field( - default="medium", description="Task priority (low, medium, high)" - ) - metadata: Dict[str, Any] = Field( - default_factory=dict, description="Additional task metadata" - ) - - -class HumanResponseRequest(BaseModel): - response: str = Field(..., description="Human response to agent question") - - -class FileOperationApprovalRequest(BaseModel): - approved: bool = Field(..., description="Whether to approve the file operations") - reason: Optional[str] = Field( - None, description="Optional reason for approval/rejection" - ) - - -class ClaudeSessionInputRequest(BaseModel): - input: str = Field(..., description="Input to send to Claude SDK session") - - -class CoordinationRequest(BaseModel): - coordination_mode: str = Field( - default="collaborative", - description="Coordination mode (sequential, parallel, hierarchical, collaborative)", - ) - required_agents: Optional[List[str]] = Field( - None, description="Specific agents to include" - ) - required_skills: Optional[List[str]] = Field( - None, description="Required skills for the task" - ) - - -class AgentCommunicationRequest(BaseModel): - message_type: str = Field( - default="notification", - description="Message type (request, response, notification, question)", - ) - content: str = Field(..., description="Message content") - metadata: Dict[str, Any] = Field( - default_factory=dict, description="Additional metadata" - ) - - -class MCPToolRequest(BaseModel): - tool_name: str = Field(..., description="Name of the MCP tool to call") - arguments: Dict[str, Any] = Field( - default_factory=dict, description="Tool arguments" - ) - - -class AgentMCPSetupRequest(BaseModel): - task_id: str = Field(..., description="Task ID for MCP setup") - session_id: Optional[str] = Field(None, description="Optional session ID") - - -class ConversationCreateRequest(BaseModel): - title: str = "New Conversation" - initial_message: Optional[str] = None - context: Optional[Dict[str, Any]] = None - - -class ConversationMessage(BaseModel): - role: str # 'user' or 'agent' - content: str - metadata: Optional[Dict[str, Any]] = None - - -class ChatMessageRequest(BaseModel): - content: str - metadata: Optional[Dict[str, Any]] = None - - -# Model Configuration Models -class ProviderCredentialsRequest(BaseModel): - provider: str = Field( - ..., description="Model provider (anthropic, openai, google, etc.)" - ) - api_key: str = Field(..., description="API key for the provider") - endpoint_url: Optional[str] = Field(None, description="Custom endpoint URL") - additional_config: Dict[str, Any] = Field( - default_factory=dict, description="Additional provider configuration" - ) - - -class AgentModelConfigRequest(BaseModel): - primary_model: str = Field(..., description="Primary model ID") - fallback_models: List[str] = Field( - default_factory=list, description="Fallback model IDs" - ) - temperature: float = Field( - default=0.7, ge=0.0, le=2.0, description="Model temperature" - ) - max_tokens: int = Field( - default=4096, ge=1, le=200000, description="Maximum output tokens" - ) - top_p: float = Field(default=1.0, ge=0.0, le=1.0, description="Top-p sampling") - frequency_penalty: float = Field( - default=0.0, ge=-2.0, le=2.0, description="Frequency penalty" - ) - presence_penalty: float = Field( - default=0.0, ge=-2.0, le=2.0, description="Presence penalty" - ) - custom_instructions: str = Field( - default="", description="Custom instructions for the agent" - ) - use_function_calling: bool = Field( - default=True, description="Enable function calling" - ) - streaming_enabled: bool = Field( - default=True, description="Enable response streaming" - ) - cost_limit_per_task: Optional[float] = Field( - None, ge=0.0, description="Cost limit per task in USD" - ) - - -class TaskCostEstimateRequest(BaseModel): - task_description: str = Field(..., description="Description of the task") - estimated_complexity: str = Field( - default="medium", - description="Estimated complexity (low, medium, high, very_high)", - ) - - -# Response models -class AgentResponse(BaseModel): - agent_id: str - status: str - agent: Dict[str, Any] - - -class TaskResponse(BaseModel): - task_id: str - status: str - - -class CoordinationResponse(BaseModel): - task_id: str - coordination_session_id: Optional[str] = None - status: str - coordination_mode: Optional[str] = None - message: Optional[str] = None - - -# Goals Management API Models -class GoalCreateRequest(BaseModel): - title: str = Field(..., description="Goal title") - description: str = Field(..., description="Goal description") - goal_type: str = Field( - default="business", - description="Goal type (business, technical, growth, operational)", - ) - target_value: Optional[Decimal] = Field( - None, description="Target value (e.g., 100000 for $100K)" - ) - target_unit: Optional[str] = Field( - None, description="Target unit (e.g., 'USD', 'users', '%')" - ) - target_deadline: Optional[date] = Field(None, description="Target completion date") - priority_level: int = Field( - default=5, ge=1, le=10, description="Priority level (1-10)" - ) - success_criteria: Optional[Dict[str, Any]] = Field( - default=None, description="Success criteria" - ) - assigned_teams: Optional[List[str]] = Field( - default=None, description="Assigned team IDs" - ) - goal_owner_agent_id: Optional[str] = Field(None, description="Goal owner agent ID") - stakeholder_agents: Optional[List[str]] = Field( - default=None, description="Stakeholder agent IDs" - ) - tags: Optional[List[str]] = Field(default=None, description="Goal tags") - metadata: Optional[Dict[str, Any]] = Field( - default=None, description="Additional metadata" - ) - - -class GoalUpdateRequest(BaseModel): - progress_percentage: Optional[Decimal] = Field( - None, ge=0, le=100, description="Progress percentage" - ) - current_value: Optional[Decimal] = Field(None, description="Current value") - completion_confidence: Optional[Decimal] = Field( - None, ge=0, le=1, description="Completion confidence" - ) - notes: Optional[str] = Field(None, description="Progress notes") - - -class MilestoneCreateRequest(BaseModel): - title: str = Field(..., description="Milestone title") - description: str = Field(..., description="Milestone description") - target_date: date = Field(..., description="Target completion date") - milestone_type: str = Field(default="deliverable", description="Milestone type") - success_criteria: Optional[Dict[str, Any]] = Field( - default=None, description="Success criteria" - ) - deliverables: Optional[List[Dict[str, Any]]] = Field( - default=None, description="Expected deliverables" - ) - dependencies: Optional[List[Dict[str, Any]]] = Field( - default=None, description="Dependencies" - ) - assigned_teams: Optional[List[str]] = Field( - default=None, description="Assigned teams" - ) - responsible_agent_id: Optional[str] = Field(None, description="Responsible agent") - priority_level: int = Field(default=5, ge=1, le=10, description="Priority level") - weight_in_goal: Optional[Decimal] = Field( - None, ge=0, le=100, description="Weight in goal (%)" - ) - - -class TaskFromMilestoneRequest(BaseModel): - title: str = Field(..., description="Task title") - description: str = Field(..., description="Task description") - task_type: str = Field(default="development", description="Task type") - complexity_level: str = Field(default="medium", description="Complexity level") - estimated_hours: Optional[Decimal] = Field(None, description="Estimated hours") - due_date: Optional[date] = Field(None, description="Due date") - assigned_team_id: Optional[str] = Field(None, description="Assigned team ID") - assigned_agent_id: Optional[str] = Field(None, description="Assigned agent ID") - priority: int = Field(default=5, ge=1, le=10, description="Priority") - requirements: Optional[Dict[str, Any]] = Field( - default=None, description="Requirements" - ) - acceptance_criteria: Optional[List[Dict[str, Any]]] = Field( - default=None, description="Acceptance criteria" - ) - dependencies: Optional[List[Dict[str, Any]]] = Field( - default=None, description="Dependencies" - ) - - -class GoalConversationCreateRequest(BaseModel): - conversation_type: str = Field(default="planning", description="Conversation type") - conversation_title: str = Field(..., description="Conversation title") - initial_context: Optional[Dict[str, Any]] = Field( - default=None, description="Initial context" - ) - participants: Optional[List[Dict[str, Any]]] = Field( - default=None, description="Participants" - ) - - -class ConversationMessageRequest(BaseModel): - message_type: str = Field(default="human", description="Message type") - sender_name: str = Field(..., description="Sender name") - content: str = Field(..., description="Message content") - metadata: Optional[Dict[str, Any]] = Field( - default=None, description="Message metadata" - ) - references: Optional[List[str]] = Field( - default=None, description="Referenced message IDs" - ) - - -class ProgressUpdateRequest(BaseModel): - progress_percentage: Optional[Decimal] = Field( - None, ge=0, le=100, description="Progress percentage" - ) - current_value: Optional[Decimal] = Field(None, description="Current value") - milestone_id: Optional[str] = Field(None, description="Associated milestone ID") - notes: Optional[str] = Field(None, description="Progress notes") - confidence_score: Optional[Decimal] = Field( - None, ge=0, le=1, description="Confidence score" - ) - trigger_alerts: bool = Field(default=True, description="Whether to trigger alerts") - - -@asynccontextmanager -async def lifespan(app: FastAPI): - # Startup - database_url = os.getenv( - "DATABASE_URL", "postgresql://postgres:password@postgres:5432/ai_context" - ) - - app.state.agent_manager = AgentManager(database_url) - app.state.task_queue = TaskQueue() - app.state.context_service = ContextService() - - # Initialize sandbox manager - app.state.sandbox_manager = AgentSandboxManager(database_url) - await app.state.sandbox_manager.start() - - # Start WebSocket manager background cleanup task - await websocket_manager.start() - - # Initialize task execution engine - app.state.task_execution_engine = TaskExecutionEngine(app.state.sandbox_manager) - await app.state.task_execution_engine.start() - - # Initialize multi-agent coordinator - from .multi_agent_coordinator import integrate_multi_agent_coordination - - app.state.multi_agent_coordinator = integrate_multi_agent_coordination( - app.state.task_execution_engine - ) - await app.state.multi_agent_coordinator.start() - - # Initialize knowledge management system - try: - from .context_enhancement_service import ContextEnhancementService - from .knowledge_notification_service import KnowledgeNotificationService - from .knowledge_propagation_engine import KnowledgePropagationEngine - from .organization_rag_manager import OrganizationRAGManager - from .task_knowledge_extractor import TaskKnowledgeExtractor - from .team_knowledge_manager import TeamKnowledgeManager - - app.state.org_rag_manager = OrganizationRAGManager(database_url) - await app.state.org_rag_manager.initialize() - - app.state.team_knowledge_manager = TeamKnowledgeManager(database_url) - await app.state.team_knowledge_manager.initialize() - - app.state.knowledge_propagation_engine = KnowledgePropagationEngine( - database_url, app.state.org_rag_manager, app.state.team_knowledge_manager - ) - await app.state.knowledge_propagation_engine.initialize() - - app.state.notification_service = KnowledgeNotificationService(database_url) - await app.state.notification_service.initialize() - - app.state.task_knowledge_extractor = TaskKnowledgeExtractor( - database_url, - app.state.org_rag_manager, - app.state.team_knowledge_manager, - app.state.knowledge_propagation_engine, - ) - await app.state.task_knowledge_extractor.initialize() - - app.state.context_enhancement_service = ContextEnhancementService( - database_url, app.state.org_rag_manager, app.state.team_knowledge_manager - ) - await app.state.context_enhancement_service.initialize() - - # Initialize knowledge analytics service - from .knowledge_analytics_service import KnowledgeAnalyticsService - - app.state.knowledge_analytics_service = KnowledgeAnalyticsService(database_url) - await app.state.knowledge_analytics_service.initialize() - - logger.info("Knowledge management system initialized successfully") - - except Exception as e: - logger.warning(f"Failed to initialize knowledge management system: {e}") - - # Initialize goals management system - try: - from .goal_conversation_service import GoalConversationService - from .goal_tracking_service import GoalTrackingService - from .goals_management_service import GoalsManagementService - from .milestone_task_engine import MilestoneTaskEngine - - app.state.goals_service = GoalsManagementService(database_url) - await app.state.goals_service.initialize() - - app.state.milestone_task_engine = MilestoneTaskEngine(database_url) - await app.state.milestone_task_engine.initialize() - - app.state.goal_conversation_service = GoalConversationService(database_url) - await app.state.goal_conversation_service.initialize() - - app.state.goal_tracking_service = GoalTrackingService(database_url) - await app.state.goal_tracking_service.initialize() - - logger.info("Goals management system initialized successfully") - - except Exception as e: - logger.warning(f"Failed to initialize goals management system: {e}") - - # Connect components - app.state.task_queue.set_task_execution_engine(app.state.task_execution_engine) - await app.state.agent_manager.set_sandbox_manager(app.state.sandbox_manager) - - # Initialize IzzyAI CEO on startup - try: - await app.state.agent_manager.create_agent( - name="IzzyAI", - role="Digital CEO", - type="executive", - config={ - "model": "claude-sonnet-4-20250514", - "temperature": 0.7, - "tools": [ - "strategic_planning", - "resource_allocation", - "team_management", - ], - }, - ) - except Exception as e: - print(f"Warning: Could not create IzzyAI CEO: {e}") - - yield - - # Shutdown - await app.state.multi_agent_coordinator.stop() - await app.state.task_execution_engine.stop() - await app.state.sandbox_manager.stop() - await app.state.agent_manager.shutdown_all() - await app.state.task_queue.close() - - # Shutdown knowledge management services - try: - if hasattr(app.state, "knowledge_analytics_service"): - await app.state.knowledge_analytics_service.close() - if hasattr(app.state, "context_enhancement_service"): - await app.state.context_enhancement_service.close() - if hasattr(app.state, "task_knowledge_extractor"): - await app.state.task_knowledge_extractor.close() - if hasattr(app.state, "notification_service"): - await app.state.notification_service.close() - if hasattr(app.state, "knowledge_propagation_engine"): - await app.state.knowledge_propagation_engine.close() - if hasattr(app.state, "team_knowledge_manager"): - await app.state.team_knowledge_manager.close() - if hasattr(app.state, "org_rag_manager"): - await app.state.org_rag_manager.close() - logger.info("Knowledge management system shutdown complete") - except Exception as e: - logger.error(f"Error shutting down knowledge management system: {e}") - - # Shutdown goals management services - try: - if hasattr(app.state, "goal_tracking_service"): - await app.state.goal_tracking_service.close() - if hasattr(app.state, "goal_conversation_service"): - await app.state.goal_conversation_service.close() - if hasattr(app.state, "milestone_task_engine"): - await app.state.milestone_task_engine.close() - if hasattr(app.state, "goals_service"): - await app.state.goals_service.close() - logger.info("Goals management system shutdown complete") - except Exception as e: - logger.error(f"Error shutting down goals management system: {e}") - - -app = FastAPI( - title="FuzeAgent Orchestrator API", - description=""" - ## FuzeAgent AI Team Orchestration Platform - - A comprehensive platform for autonomous AI development teams that enables: - - ### 🤖 Autonomous Agent Execution - - **Claude SDK Integration**: Interactive AI development with real-time conversation streaming - - **File Operations Engine**: Safe code changes with human approval workflows - - **Multi-Agent Coordination**: Complex task decomposition and agent collaboration - - ### 🏗️ Core Features - - **Agent Management**: Create, configure, and manage AI development agents - - **Task Orchestration**: Assign and monitor complex development tasks - - **Real-time Monitoring**: WebSocket streaming for live progress updates - - **Human-in-the-Loop**: Seamless approval workflows for critical decisions - - ### 🔗 Integration Capabilities - - **MCP (Model Context Protocol)**: Organizational context for AI agents - - **Git Workflow Management**: Automated repository operations - - **Sandbox Environments**: Isolated development containers - - **Database Integration**: PostgreSQL for persistent storage - - ### 📡 API Categories - - **Agent Management**: Create and manage AI agents - - **Task Execution**: Autonomous task processing - - **File Operations**: Code change management - - **Multi-Agent Coordination**: Team collaboration - - **Real-time Communication**: WebSocket endpoints - - **MCP Integration**: Organizational context tools - - **Goals Management**: Organizational goals, milestones, and task planning - - **Knowledge Management**: RAG system and organizational memory - - **Version**: 2.0.0 (Autonomous Execution) - """, - version="2.0.0", - lifespan=lifespan, - docs_url="/docs", - redoc_url="/redoc", - openapi_tags=[ - {"name": "health", "description": "Health check and system status endpoints"}, - { - "name": "agents", - "description": "AI agent creation, management, and status monitoring", - }, - {"name": "tasks", "description": "Task assignment, execution, and monitoring"}, - { - "name": "autonomous-execution", - "description": "Autonomous task execution with Claude SDK integration", - }, - { - "name": "file-operations", - "description": "File system operations and code change management", - }, - { - "name": "multi-agent-coordination", - "description": "Multi-agent collaboration and task coordination", - }, - { - "name": "real-time", - "description": "WebSocket endpoints for real-time updates", - }, - { - "name": "human-in-loop", - "description": "Human approval workflows and interaction handling", - }, - { - "name": "mcp-integration", - "description": "Model Context Protocol tools and resources", - }, - {"name": "sandboxes", "description": "Sandbox environment management"}, - {"name": "context", "description": "Agent memory and context management"}, - { - "name": "model-configuration", - "description": "AI model configuration and API key management", - }, - { - "name": "knowledge-management", - "description": "Hierarchical knowledge management, RAG, and intelligent notifications", - }, - ], -) - -app.add_middleware( - CORSMiddleware, - allow_origins=[ - "http://localhost:3000", - "http://localhost:3031", - "http://localhost", - ], - allow_credentials=True, - allow_methods=["*"], - allow_headers=["*"], -) - -# Include hierarchy router for organizational visualization -app.include_router(hierarchy_router) - - -# Health check endpoint -@app.get( - "/health", - tags=["health"], - summary="Health Check", - description="Check the health status of the FuzeAgent orchestrator service", - response_description="Service health status", -) -async def health_check(): - """ - Health check endpoint that returns the current status of the orchestrator service. - - Returns: - dict: Service health status and basic information - """ - return { - "status": "healthy", - "service": "orchestrator", - "version": "2.0.0", - "features": { - "autonomous_execution": True, - "multi_agent_coordination": True, - "file_operations": True, - "mcp_integration": True, - "real_time_streaming": True, - }, - # Whether GET /openapi.yaml can answer. An image built without its - # contract is DEGRADED, not dead — the probe still passes (no restart - # can conjure a file the image lacks) but the condition is visible to - # anything that looks, instead of surfacing only as a 503 later. - "openapi": "loaded" if _openapi_document() is not None else "unavailable", - } - - -# --------------------------------------------------------------------------- -# The contract, SERVED. -# -# contracts/openapi.yaml describes this orchestrator's real HTTP surface, with -# the curated descriptions and the irreversibility guidance that -# mcp/tools.overrides.yaml narrows. Committing it is not the same as publishing -# it: consumers — the MCP gateway among them — discover the surface over HTTP. -# -# This is NOT /openapi.json. FastAPI generates that from the code at import -# time; it is accurate about shapes and says nothing about which operations -# dispatch an agent that cannot be recalled. Both are served. This one is the -# contract. -# -# The document is read from the IMAGE, never from a mount, so what this endpoint -# publishes is always the contract this build was compiled against. -# --------------------------------------------------------------------------- -_ORCH_DIR = os.path.dirname(os.path.abspath(__file__)) -_OPENAPI_CANDIDATES = [ - p - for p in [ - os.getenv("OPENAPI_SPEC_PATH"), - os.path.join(_ORCH_DIR, "contracts", "openapi.yaml"), - os.path.join(_ORCH_DIR, "..", "..", "contracts", "openapi.yaml"), - ] - if p -] - - -def _openapi_document(): - """Return the OpenAPI document text, or None when the image lacks it.""" - for path in _OPENAPI_CANDIDATES: - try: - with open(path, "r", encoding="utf-8") as fh: - return fh.read() - except OSError: - continue - return None - - -@app.get( - "/openapi.yaml", - tags=["health"], - summary="This OpenAPI Document", - description=( - "Serve contracts/openapi.yaml — the curated contract, as distinct from " - "FastAPI's auto-generated /openapi.json." - ), - include_in_schema=False, -) -async def get_openapi_document(): - doc = _openapi_document() - if doc is None: - logger.error("OpenAPI document not found; tried %s", _OPENAPI_CANDIDATES) - # 503, not 500 and not a crash: the service is otherwise functional and - # no restart can produce a spec the image does not contain. - raise HTTPException( - status_code=503, - detail=( - "openapi_document_unavailable: this image was built without " - "contracts/openapi.yaml. Rebuild with the repo root as the Docker " - "context so the contract is copied in." - ), - ) - return Response(content=doc, media_type="application/yaml") - - -# WebSocket for real-time updates -@app.websocket("/ws") -async def websocket_endpoint(websocket: WebSocket): - await websocket.accept() - try: - while True: - # Send agent updates to UI - updates = await app.state.agent_manager.get_updates() - await websocket.send_json(updates) - await asyncio.sleep(1) - except Exception as e: - print(f"WebSocket error: {e}") - finally: - await websocket.close() - - -# WebSocket for task execution updates -@app.websocket("/ws/tasks/{task_id}") -async def task_websocket_endpoint(websocket: WebSocket, task_id: str): - """WebSocket endpoint for real-time task execution updates""" - await websocket.accept() - try: - while True: - # Get task execution status - try: - status = await app.state.task_queue.get_execution_status(task_id) - await websocket.send_json( - {"type": "status_update", "task_id": task_id, "data": status} - ) - - # If task is completed or failed, send final update and close - if status.get("status") in ["completed", "failed", "cancelled"]: - await websocket.send_json( - { - "type": "task_finished", - "task_id": task_id, - "final_status": status.get("status"), - } - ) - break - - except Exception as e: - await websocket.send_json( - {"type": "error", "task_id": task_id, "error": str(e)} - ) - - await asyncio.sleep(2) # Update every 2 seconds - - except Exception as e: - print(f"Task WebSocket error for {task_id}: {e}") - finally: - await websocket.close() - - -# WebSocket for real-time Claude SDK conversation streaming -@app.websocket("/ws/tasks/{task_id}/conversation") -async def conversation_websocket_endpoint(websocket: WebSocket, task_id: str): - """WebSocket endpoint for real-time Claude SDK conversation streaming""" - await websocket.accept() - try: - # Get execution context - execution = app.state.task_execution_engine.active_executions.get(task_id) - if not execution: - await websocket.send_json( - {"type": "error", "message": f"Task {task_id} not found or not active"} - ) - await websocket.close() - return - - # Wait for Claude SDK session to be available - while not execution.claude_session_id and execution.status not in [ - "completed", - "failed", - "cancelled", - ]: - await asyncio.sleep(1) - - if not execution.claude_session_id: - await websocket.send_json( - { - "type": "error", - "message": "No active Claude SDK session for this task", - } - ) - await websocket.close() - return - - # Stream Claude SDK output - claude_sdk_manager = execution.claude_sdk_manager - if claude_sdk_manager: - try: - async for output_chunk in claude_sdk_manager.stream_output( - execution.claude_session_id - ): - await websocket.send_json( - { - "type": "claude_output", - "task_id": task_id, - "content": output_chunk, - "timestamp": datetime.now().isoformat(), - } - ) - - # Session ended - await websocket.send_json( - { - "type": "conversation_ended", - "task_id": task_id, - "timestamp": datetime.now().isoformat(), - } - ) - - except Exception as e: - await websocket.send_json( - { - "type": "error", - "message": f"Error streaming conversation: {str(e)}", - } - ) - - except Exception as e: - print(f"Conversation WebSocket error for {task_id}: {e}") - finally: - await websocket.close() - - -# WebSocket for file operations streaming -@app.websocket("/ws/tasks/{task_id}/file-operations") -async def file_operations_websocket_endpoint(websocket: WebSocket, task_id: str): - """WebSocket endpoint for real-time file operations updates""" - await websocket.accept() - try: - # Get execution context - execution = app.state.task_execution_engine.active_executions.get(task_id) - if not execution: - await websocket.send_json( - {"type": "error", "message": f"Task {task_id} not found or not active"} - ) - await websocket.close() - return - - file_ops_engine = execution.file_operations_engine - if not file_ops_engine: - await websocket.send_json( - { - "type": "error", - "message": "No file operations engine available for this task", - } - ) - await websocket.close() - return - - last_batch_count = 0 - - while execution.status not in ["completed", "failed", "cancelled"]: - try: - # Get pending operations - pending_operations = file_ops_engine.get_pending_operations() - applied_operations = file_ops_engine.get_applied_operations() - - current_batch_count = len(pending_operations) + len(applied_operations) - - # Send updates if there are new operations - if current_batch_count > last_batch_count: - # Send pending operations - for batch in pending_operations: - # Get diff preview - diffs = await file_ops_engine.get_file_diff_preview( - batch.batch_id - ) - - await websocket.send_json( - { - "type": "pending_operations", - "task_id": task_id, - "batch_id": batch.batch_id, - "description": batch.description, - "requires_approval": batch.requires_approval, - "operations_count": len(batch.operations), - "file_diffs": diffs, - "timestamp": batch.created_at.isoformat(), - } - ) - - # Send applied operations - for batch in applied_operations: - await websocket.send_json( - { - "type": "applied_operations", - "task_id": task_id, - "batch_id": batch.batch_id, - "description": batch.description, - "operations_count": len(batch.operations), - "applied_at": ( - batch.applied_at.isoformat() - if batch.applied_at - else None - ), - "timestamp": batch.created_at.isoformat(), - } - ) - - last_batch_count = current_batch_count - - await asyncio.sleep(1) # Check every second - - except Exception as e: - await websocket.send_json( - { - "type": "error", - "message": f"Error getting file operations: {str(e)}", - } - ) - - # Task completed - await websocket.send_json( - { - "type": "task_completed", - "task_id": task_id, - "final_status": execution.status.value, - "timestamp": datetime.now().isoformat(), - } - ) - - except Exception as e: - print(f"File operations WebSocket error for {task_id}: {e}") - finally: - await websocket.close() - - -# Agent Management Endpoints -@app.post( - "/agents", - tags=["agents"], - summary="Create AI Agent", - description="Create a new AI agent with repository and sandbox settings", - response_model=AgentResponse, -) -async def create_agent(agent_config: AgentCreateRequest, _auth=Depends(require_auth)): - """Create a new AI agent with repository and sandbox settings""" - try: - agent = await app.state.agent_manager.create_agent(**agent_config) - return { - "agent_id": agent.id, - "status": "created", - "agent": { - "id": agent.id, - "name": agent_config.get("name"), - "role": agent_config.get("role"), - "type": agent_config.get("type"), - "repository_settings": agent_config.get("repository_settings", {}), - "sandbox_settings": agent_config.get("sandbox_settings", {}), - "created_at": ( - agent.created_at if hasattr(agent, "created_at") else None - ), - }, - } - except Exception as e: - raise HTTPException(status_code=400, detail=f"Failed to create agent: {str(e)}") - - -@app.get( - "/agents", - tags=["agents"], - summary="List All Agents", - description="Get a list of all AI agents and their current status", -) -async def list_agents(): - """List all agents and their status""" - return await app.state.agent_manager.list_agents() - - -@app.post( - "/agents/{agent_id}/tasks", - tags=["tasks"], - summary="Assign Task to Agent", - description="Assign a specific task to an AI agent", - response_model=TaskResponse, -) -async def assign_task( - agent_id: str = Path(..., description="Agent ID"), - task: TaskCreateRequest = Body(...), - _auth=Depends(require_auth), -): - """Assign a task to an agent""" - task_id = await app.state.task_queue.assign_task(agent_id, task) - return {"task_id": task_id, "status": "assigned"} - - -@app.get("/agents/{agent_id}/status") -async def get_agent_status(agent_id: str): - """Get detailed agent status""" - return await app.state.agent_manager.get_agent_status(agent_id) - - -@app.get("/agents/{agent_id}/tasks") -async def get_agent_tasks(agent_id: str): - """Get tasks assigned to an agent""" - try: - # This would normally query the database for tasks assigned to the agent - # For now, return mock data - return [ - { - "id": "1", - "title": "Strategic Planning Q4 2025", - "description": "Develop comprehensive strategic plan for Q4 2025 expansion", - "status": "completed", - "priority": "high", - "created_at": "2025-08-05T09:00:00Z", - "completed_at": "2025-08-05T17:30:00Z", - }, - { - "id": "2", - "title": "Team Performance Review", - "description": "Conduct quarterly performance review for all team leads", - "status": "in_progress", - "priority": "medium", - "created_at": "2025-08-06T08:00:00Z", - }, - ] - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get agent tasks: {str(e)}" - ) - - -@app.get("/teams") -async def list_teams(): - """List all teams""" - try: - # This would normally query the database for teams - # For now, return mock data - return [ - { - "id": "1", - "name": "Executive Team", - "description": "Strategic leadership and decision making", - }, - { - "id": "2", - "name": "Development Team", - "description": "Frontend, backend, and full-stack development", - }, - { - "id": "3", - "name": "Quality Assurance", - "description": "Testing, quality control, and bug detection", - }, - { - "id": "4", - "name": "DevOps Team", - "description": "Infrastructure, deployment, and system operations", - }, - { - "id": "5", - "name": "Business Team", - "description": "Marketing, sales, and customer relations", - }, - ] - except Exception as e: - raise HTTPException(status_code=500, detail=f"Failed to list teams: {str(e)}") - - -@app.get("/agent-templates") -async def list_agent_templates(): - """List available agent templates""" - try: - return [ - { - "id": "react_developer", - "name": "React Developer", - "description": "Frontend developer specialized in React and TypeScript", - "type": "developer", - "defaultConfig": { - "model": "claude-sonnet-4-20250514", - "temperature": 0.7, - "tools": ["code_generation", "code_review", "debugging", "testing"], - "goal": "Build responsive and performant React applications", - "backstory": "Experienced frontend developer with deep knowledge of React ecosystem", - }, - }, - { - "id": "python_developer", - "name": "Python Developer", - "description": "Backend developer specialized in Python and FastAPI", - "type": "developer", - "defaultConfig": { - "model": "claude-sonnet-4-20250514", - "temperature": 0.7, - "tools": [ - "code_generation", - "api_development", - "database_design", - "testing", - ], - "goal": "Develop robust and scalable backend systems", - "backstory": "Senior Python developer with expertise in FastAPI and databases", - }, - }, - { - "id": "qa_engineer", - "name": "QA Engineer", - "description": "Quality assurance engineer focused on testing and automation", - "type": "qa", - "defaultConfig": { - "model": "claude-sonnet-4-20250514", - "temperature": 0.6, - "tools": [ - "test_automation", - "bug_reporting", - "quality_analysis", - "performance_testing", - ], - "goal": "Ensure high quality and reliability of software products", - "backstory": "Experienced QA engineer with expertise in automated testing frameworks", - }, - }, - { - "id": "devops_engineer", - "name": "DevOps Engineer", - "description": "Infrastructure and deployment specialist", - "type": "devops", - "defaultConfig": { - "model": "claude-sonnet-4-20250514", - "temperature": 0.5, - "tools": [ - "infrastructure_management", - "deployment", - "monitoring", - "security", - ], - "goal": "Maintain reliable and scalable infrastructure", - "backstory": "DevOps engineer with expertise in cloud platforms and CI/CD", - }, - }, - ] - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to list agent templates: {str(e)}" - ) - - -@app.get("/tasks") -async def list_tasks(): - """List all tasks""" - return await app.state.task_queue.list_tasks() - - -@app.get("/tasks/{task_id}") -async def get_task(task_id: str): - """Get task details""" - return await app.state.task_queue.get_task(task_id) - - -# Autonomous Execution Endpoints -@app.post("/agents/from-template") -async def create_agent_from_template(request: dict): - """Create agent from template with repository settings""" - try: - # Extract template data - template_id = request.get("template_id") - name = request.get("name") - team_id = request.get("team_id") - overrides = request.get("overrides", {}) - - # Get template configuration - template_config = await app.state.agent_manager.get_template_config(template_id) - if not template_config: - raise HTTPException( - status_code=404, detail=f"Template {template_id} not found" - ) - - # Build agent configuration - agent_config = { - "name": name, - "role": template_config.get("role", template_id.replace("_", " ").title()), - "type": template_config.get("type", "specialized"), - "template_id": template_id, - "team_id": team_id, - "config": {**template_config.get("config", {}), **overrides}, - "repository_settings": request.get("repository_settings", {}), - "sandbox_settings": { - "base_image": f"fuzeagent/dev-{template_id.split('_')[0]}:latest", - "resource_limits": template_config.get( - "resource_limits", {"memory": "2Gi", "cpu": "1.0", "disk": "10Gi"} - ), - "auto_cleanup": "24h", - }, - } - - # Create agent - agent = await app.state.agent_manager.create_agent(**agent_config) - - return { - "agent_id": agent.id, - "status": "created", - "agent": agent_config, - "template_id": template_id, - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to create agent from template: {str(e)}" - ) - - -@app.get("/templates") -async def get_agent_templates(): - """Get available agent templates""" - return await app.state.agent_manager.get_available_templates() - - -@app.post( - "/tasks/{task_id}/execute", - tags=["autonomous-execution"], - summary="Start Autonomous Task Execution", - description="Begin autonomous execution of a task using Claude SDK integration", - response_model=TaskResponse, -) -async def start_task_execution( - task_id: str = Path(..., description="Task ID to execute") -): - """Start autonomous execution of a task""" - try: - # This will be handled by the TaskExecutionEngine - result = await app.state.task_queue.start_autonomous_execution(task_id) - return {"task_id": task_id, "status": "execution_started", "result": result} - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to start task execution: {str(e)}" - ) - - -@app.get("/tasks/{task_id}/status") -async def get_task_execution_status(task_id: str): - """Get detailed task execution status""" - try: - status = await app.state.task_queue.get_execution_status(task_id) - return status - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get task status: {str(e)}" - ) - - -@app.get("/tasks/{task_id}/iterations") -async def get_task_iterations(task_id: str): - """Get task iteration history""" - try: - iterations = await app.state.task_queue.get_task_iterations(task_id) - return {"task_id": task_id, "iterations": iterations} - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get task iterations: {str(e)}" - ) - - -@app.get("/agents/{agent_id}/sandbox") -async def get_agent_sandbox(agent_id: str): - """Get agent sandbox information""" - try: - sandbox_info = await app.state.agent_manager.get_agent_sandbox(agent_id) - return sandbox_info - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get agent sandbox: {str(e)}" - ) - - -# Additional endpoints for UI support -@app.put("/tasks/{task_id}") -async def update_task(task_id: str, update_data: dict): - """Update task status and result""" - await app.state.task_queue.update_task_status( - task_id=task_id, - status=update_data.get("status"), - result=update_data.get("result"), - ) - return {"status": "updated"} - - -@app.post("/context/interactions") -async def store_interaction(interaction_data: dict): - """Store agent interaction""" - interaction_id = await app.state.context_service.store_interaction( - agent_id=interaction_data.get("agent_id"), - content=interaction_data.get("content"), - metadata=interaction_data.get("metadata", {}), - ) - return {"interaction_id": interaction_id} - - -@app.get("/context") -async def get_context(query: str, agent_id: str = None): - """Get relevant context for a query""" - context = await app.state.context_service.get_context(query, agent_id) - return context - - -@app.get("/agents/{agent_id}/memory") -async def get_agent_memory(agent_id: str, limit: int = 10): - """Get agent memory""" - memory = await app.state.context_service.get_agent_memory(agent_id, limit) - return memory - - -# Agent Conversation Endpoints -@app.get( - "/agents/{agent_id}/conversations", - tags=["conversations"], - summary="Get Agent Conversations", - description="Get all conversations for a specific agent", -) -async def get_agent_conversations(agent_id: str): - """Get all conversations for a specific agent""" - try: - async with get_db_connection() as conn: - conversations = await conn.fetch( - """ - SELECT cs.*, COUNT(ac.id) as message_count, - (SELECT content FROM agent_conversations - WHERE session_id = cs.id - ORDER BY created_at DESC LIMIT 1) as last_message - FROM chat_sessions cs - LEFT JOIN agent_conversations ac ON cs.id = ac.session_id - WHERE cs.agent_id = $1 - GROUP BY cs.id - ORDER BY cs.last_activity DESC - """, - agent_id, - ) - - return [dict(row) for row in conversations] - - except Exception as e: - logger.error(f"Error getting agent conversations: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/agents/{agent_id}/conversations", - tags=["conversations"], - summary="Create New Agent Conversation", - description="Create a new conversation with an agent", -) -async def create_agent_conversation(agent_id: str, request: ConversationCreateRequest): - """Create a new conversation with an agent""" - try: - async with get_db_connection() as conn: - # Create new chat session - session_id = await conn.fetchval( - """ - INSERT INTO chat_sessions (agent_id, session_name, context, status) - VALUES ($1, $2, $3, 'active') - RETURNING id - """, - agent_id, - request.title, - request.context or {}, - ) - - # Add initial message if provided - if request.initial_message: - await conn.execute( - """ - INSERT INTO agent_conversations (session_id, agent_id, message_type, content) - VALUES ($1, $2, 'system', $3) - """, - session_id, - agent_id, - request.initial_message, - ) - - # Get the created conversation - conversation = await conn.fetchrow( - """ - SELECT * FROM chat_sessions WHERE id = $1 - """, - session_id, - ) - - return dict(conversation) - - except Exception as e: - logger.error(f"Error creating agent conversation: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/agents/{agent_id}/conversations/{conversation_id}/messages", - tags=["conversations"], - summary="Get Conversation Messages", - description="Get all messages in a conversation", -) -async def get_conversation_messages(agent_id: str, conversation_id: str): - """Get all messages in a conversation""" - try: - async with get_db_connection() as conn: - messages = await conn.fetch( - """ - SELECT * FROM agent_conversations - WHERE session_id = $1 AND agent_id = $2 - ORDER BY created_at ASC - """, - conversation_id, - agent_id, - ) - - return [dict(row) for row in messages] - - except Exception as e: - logger.error(f"Error getting conversation messages: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/agents/{agent_id}/conversations/{conversation_id}/messages", - tags=["conversations"], - summary="Send Message to Agent", - description="Send a message to an agent in a conversation", -) -async def send_message_to_agent( - agent_id: str, conversation_id: str, request: ChatMessageRequest -): - """Send a message to an agent in a conversation""" - try: - async with get_db_connection() as conn: - # Insert user message - user_message_id = await conn.fetchval( - """ - INSERT INTO agent_conversations (session_id, agent_id, message_type, content, metadata) - VALUES ($1, $2, 'user', $3, $4) - RETURNING id - """, - conversation_id, - agent_id, - request.content, - request.metadata or {}, - ) - - # Update session last activity - await conn.execute( - """ - UPDATE chat_sessions - SET last_activity = CURRENT_TIMESTAMP, message_count = message_count + 1 - WHERE id = $1 - """, - conversation_id, - ) - - # TODO: Here we would integrate with the actual agent to generate a response - # For now, return a simple acknowledgment - - return { - "id": str(user_message_id), - "status": "sent", - "message": "Message sent to agent", - } - - except Exception as e: - logger.error(f"Error sending message to agent: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get("/knowledge/search") -async def search_knowledge(query: str, limit: int = 10): - """Search knowledge base""" - results = await app.state.context_service.search_knowledge(query, limit) - return results - - -# Human-in-the-loop endpoints -@app.post( - "/tasks/{task_id}/human-response", - tags=["human-in-loop"], - summary="Submit Human Response", - description="Submit human response to a task question or approval request", -) -async def submit_human_response( - task_id: str = Path(..., description="Task ID"), - response_data: HumanResponseRequest = Body(...), -): - """Submit human response to a task question""" - try: - response = response_data.get("response", "") - if not response: - raise HTTPException(status_code=400, detail="Response cannot be empty") - - success = await app.state.task_queue.handle_human_response(task_id, response) - - if success: - return {"status": "success", "message": "Human response submitted"} - else: - raise HTTPException( - status_code=404, - detail="Task not found or not waiting for human response", - ) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to submit human response: {str(e)}" - ) - - -@app.post("/tasks/{task_id}/cancel") -async def cancel_task_execution(task_id: str): - """Cancel autonomous execution of a task""" - try: - success = await app.state.task_queue.cancel_task_execution(task_id) - - if success: - return {"status": "cancelled", "message": "Task execution cancelled"} - else: - raise HTTPException(status_code=404, detail="Task not found or not running") - - except Exception as e: - raise HTTPException(status_code=500, detail=f"Failed to cancel task: {str(e)}") - - -@app.get("/tasks/{task_id}/messages") -async def get_task_messages(task_id: str): - """Get task messages and chat history""" - try: - # This would integrate with the HumanInTheLoopHandler when implemented - # For now, return iteration history which includes human interactions - iterations = await app.state.task_queue.get_task_iterations(task_id) - - messages = [] - for iteration in iterations: - if iteration.get("human_question"): - messages.append( - { - "type": "agent_question", - "content": iteration["human_question"], - "timestamp": iteration["started_at"], - "iteration": iteration["iteration_number"], - } - ) - - if iteration.get("human_response"): - messages.append( - { - "type": "human_response", - "content": iteration["human_response"], - "timestamp": iteration["completed_at"] - or iteration["started_at"], - "iteration": iteration["iteration_number"], - } - ) - - return {"task_id": task_id, "messages": messages} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get task messages: {str(e)}" - ) - - -# Sandbox management endpoints -@app.get("/sandboxes") -async def list_sandboxes(agent_id: str = None, status: str = None): - """List active sandboxes""" - try: - from .sandbox_manager import SandboxStatus - - sandbox_status = None - if status: - try: - sandbox_status = SandboxStatus(status) - except ValueError: - raise HTTPException(status_code=400, detail=f"Invalid status: {status}") - - sandboxes = await app.state.sandbox_manager.list_sandboxes( - agent_id=agent_id, status=sandbox_status - ) - - return { - "sandboxes": [ - { - "sandbox_id": s.sandbox_id, - "agent_id": s.agent_id, - "task_id": s.task_id, - "status": s.status.value, - "workspace_path": s.workspace_path, - "created_at": s.created_at.isoformat(), - "resource_limits": s.resource_limits, - } - for s in sandboxes - ] - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to list sandboxes: {str(e)}" - ) - - -@app.post("/sandboxes/{sandbox_id}/execute") -async def execute_command_in_sandbox(sandbox_id: str, command_data: dict): - """Execute a command in a sandbox""" - try: - command = command_data.get("command") - working_dir = command_data.get("working_dir") - - if not command: - raise HTTPException(status_code=400, detail="Command is required") - - result = await app.state.sandbox_manager.execute_command( - sandbox_id=sandbox_id, command=command, working_dir=working_dir - ) - - return result - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to execute command: {str(e)}" - ) - - -@app.delete("/sandboxes/{sandbox_id}") -async def destroy_sandbox(sandbox_id: str): - """Destroy a sandbox""" - try: - await app.state.sandbox_manager.destroy_sandbox(sandbox_id) - return {"status": "destroyed", "sandbox_id": sandbox_id} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to destroy sandbox: {str(e)}" - ) - - -# Agent registration and communication endpoints -@app.post("/agents/{agent_id}/register") -async def register_agent(agent_id: str, registration_data: dict): - """Register an agent running in a sandbox container""" - try: - # Store agent registration info - # This would typically update the agent's status and capabilities - return { - "status": "registered", - "agent_id": agent_id, - "registered_at": datetime.now().isoformat(), - } - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to register agent: {str(e)}" - ) - - -@app.get("/agents/{agent_id}/next-task") -async def get_next_task_for_agent(agent_id: str): - """Get the next task for an agent to execute""" - try: - # Find pending tasks assigned to this agent - tasks = await app.state.task_queue.get_agent_tasks(agent_id) - pending_tasks = [t for t in tasks if t.get("status") == "pending"] - - if pending_tasks: - # Return the first pending task - task = pending_tasks[0] - # Update status to 'assigned' to prevent double assignment - await app.state.task_queue.update_task_status(task["id"], "assigned") - return task - else: - # No tasks available - return None, 204 - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get next task: {str(e)}" - ) - - -@app.post("/agents/{agent_id}/error") -async def report_agent_error(agent_id: str, error_data: dict): - """Report an error from an agent""" - try: - # Log the error and update agent status - logger.error(f"Agent {agent_id} reported error: {error_data.get('error')}") - - # You might want to store this in a database or alerting system - return {"status": "error_logged", "agent_id": agent_id} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to log agent error: {str(e)}" - ) - - -# Conversation management endpoints -@app.get("/tasks/{task_id}/conversation") -async def get_task_conversation(task_id: str, iteration: int = None, limit: int = 100): - """Get conversation history for a task""" - try: - conversation_history = await app.state.task_execution_engine.conversation_manager.get_conversation_history( - task_id=task_id, iteration_number=iteration, limit=limit - ) - - return {"task_id": task_id, "conversation": conversation_history} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get conversation: {str(e)}" - ) - - -@app.get("/tasks/{task_id}/conversation/summary") -async def get_conversation_summary(task_id: str): - """Get conversation summary with statistics""" - try: - summary = await app.state.task_execution_engine.conversation_manager.get_conversation_summary( - task_id - ) - return summary - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get conversation summary: {str(e)}" - ) - - -@app.get("/tasks/{task_id}/code-generations") -async def get_task_code_generations( - task_id: str, iteration: int = None, file_type: str = None -): - """Get code generations for a task""" - try: - code_generations = await app.state.task_execution_engine.conversation_manager.get_code_generations( - task_id=task_id, iteration_number=iteration, file_type=file_type - ) - - return {"task_id": task_id, "code_generations": code_generations} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get code generations: {str(e)}" - ) - - -@app.get("/agents/{agent_id}/performance") -async def get_agent_performance(agent_id: str, hours: int = 24): - """Get agent performance metrics""" - try: - metrics = await app.state.task_execution_engine.conversation_manager.get_agent_performance_metrics( - agent_id=agent_id, time_range_hours=hours - ) - - return {"agent_id": agent_id, "time_range_hours": hours, "metrics": metrics} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get agent performance: {str(e)}" - ) - - -# File Operations Endpoints -@app.get( - "/tasks/{task_id}/file-operations", - tags=["file-operations"], - summary="Get File Operations", - description="Get file operations for a task with optional status filtering", -) -async def get_task_file_operations( - task_id: str = Path(..., description="Task ID"), - status: Optional[str] = Query( - None, description="Filter by status (pending, applied)" - ), -): - """Get file operations for a task""" - try: - execution = app.state.task_execution_engine.active_executions.get(task_id) - if not execution or not execution.file_operations_engine: - raise HTTPException( - status_code=404, detail="Task not found or no file operations available" - ) - - file_ops_engine = execution.file_operations_engine - - if status == "pending": - operations = file_ops_engine.get_pending_operations() - elif status == "applied": - operations = file_ops_engine.get_applied_operations() - else: - # Get all operations - pending = file_ops_engine.get_pending_operations() - applied = file_ops_engine.get_applied_operations() - operations = pending + applied - - # Convert to dict format - operations_data = [] - for batch in operations: - operations_data.append( - { - "batch_id": batch.batch_id, - "task_id": batch.task_id, - "agent_id": batch.agent_id, - "description": batch.description, - "requires_approval": batch.requires_approval, - "approval_status": batch.approval_status.value, - "operations_count": len(batch.operations), - "created_at": batch.created_at.isoformat(), - "applied_at": ( - batch.applied_at.isoformat() if batch.applied_at else None - ), - } - ) - - return {"task_id": task_id, "operations": operations_data} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get file operations: {str(e)}" - ) - - -@app.get("/tasks/{task_id}/file-operations/{batch_id}/preview") -async def get_file_operations_preview(task_id: str, batch_id: str): - """Get preview of file changes for a batch""" - try: - execution = app.state.task_execution_engine.active_executions.get(task_id) - if not execution or not execution.file_operations_engine: - raise HTTPException( - status_code=404, detail="Task not found or no file operations available" - ) - - file_ops_engine = execution.file_operations_engine - diffs = await file_ops_engine.get_file_diff_preview(batch_id) - - return {"task_id": task_id, "batch_id": batch_id, "file_diffs": diffs} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get file preview: {str(e)}" - ) - - -@app.post( - "/tasks/{task_id}/file-operations/{batch_id}/approve", - tags=["file-operations", "human-in-loop"], - summary="Approve File Operations", - description="Approve or reject file operations from Claude SDK", -) -async def approve_file_operations( - task_id: str = Path(..., description="Task ID"), - batch_id: str = Path(..., description="Batch ID"), - approval_data: FileOperationApprovalRequest = Body(...), -): - """Approve or reject file operations""" - try: - approved = approval_data.get("approved", False) - - execution = app.state.task_execution_engine.active_executions.get(task_id) - if not execution or not execution.file_operations_engine: - raise HTTPException( - status_code=404, detail="Task not found or no file operations available" - ) - - file_ops_engine = execution.file_operations_engine - success = await file_ops_engine.approve_operations(batch_id, approved) - - if success: - # Also notify Claude SDK if there's an active session - if execution.claude_sdk_manager and execution.claude_session_id: - await execution.claude_sdk_manager.approve_file_operations( - execution.claude_session_id, batch_id, approved - ) - - return { - "task_id": task_id, - "batch_id": batch_id, - "approved": approved, - "status": "success", - } - else: - raise HTTPException(status_code=400, detail="Failed to process approval") - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to approve file operations: {str(e)}" - ) - - -@app.post("/tasks/{task_id}/file-operations/{batch_id}/rollback") -async def rollback_file_operations(task_id: str, batch_id: str): - """Rollback applied file operations""" - try: - execution = app.state.task_execution_engine.active_executions.get(task_id) - if not execution or not execution.file_operations_engine: - raise HTTPException( - status_code=404, detail="Task not found or no file operations available" - ) - - file_ops_engine = execution.file_operations_engine - success = await file_ops_engine.rollback_operations(batch_id) - - if success: - return {"task_id": task_id, "batch_id": batch_id, "status": "rolled_back"} - else: - raise HTTPException(status_code=400, detail="Failed to rollback operations") - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to rollback file operations: {str(e)}" - ) - - -# Claude SDK Session Management Endpoints -@app.get("/tasks/{task_id}/claude-session") -async def get_claude_session_status(task_id: str): - """Get Claude SDK session status for a task""" - try: - execution = app.state.task_execution_engine.active_executions.get(task_id) - if ( - not execution - or not execution.claude_sdk_manager - or not execution.claude_session_id - ): - raise HTTPException( - status_code=404, detail="No active Claude SDK session for this task" - ) - - status = await execution.claude_sdk_manager.get_session_status( - execution.claude_session_id - ) - return status - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get Claude session status: {str(e)}" - ) - - -@app.post("/tasks/{task_id}/claude-session/input") -async def send_claude_session_input(task_id: str, input_data: dict): - """Send input to Claude SDK session""" - try: - user_input = input_data.get("input", "") - if not user_input: - raise HTTPException(status_code=400, detail="Input cannot be empty") - - execution = app.state.task_execution_engine.active_executions.get(task_id) - if ( - not execution - or not execution.claude_sdk_manager - or not execution.claude_session_id - ): - raise HTTPException( - status_code=404, detail="No active Claude SDK session for this task" - ) - - success = await execution.claude_sdk_manager.send_input( - execution.claude_session_id, user_input - ) - - if success: - return {"task_id": task_id, "status": "input_sent", "input": user_input} - else: - raise HTTPException( - status_code=400, detail="Failed to send input to Claude session" - ) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to send Claude session input: {str(e)}" - ) - - -# MCP Integration Endpoints -@app.get("/mcp/tools") -async def get_mcp_tools(): - """Get available MCP tools""" - try: - from .mcp_integration import FuzeAgentMCPServer - - mcp_server = FuzeAgentMCPServer() - tools = [ - { - "name": tool.name, - "description": tool.description, - "input_schema": tool.input_schema, - } - for tool in mcp_server.tools - ] - - return {"tools": tools} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get MCP tools: {str(e)}" - ) - - -@app.post( - "/mcp/call-tool", - tags=["mcp-integration"], - summary="Call MCP Tool", - description="Execute an MCP tool to access organizational context", -) -async def call_mcp_tool(tool_request: MCPToolRequest = Body(...)): - """Call an MCP tool""" - try: - from .mcp_integration import FuzeAgentMCPServer - - tool_name = tool_request.get("tool_name") - arguments = tool_request.get("arguments", {}) - - if not tool_name: - raise HTTPException(status_code=400, detail="tool_name is required") - - mcp_server = FuzeAgentMCPServer() - result = await mcp_server.handle_tool_call(tool_name, arguments) - - return result - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to call MCP tool: {str(e)}" - ) - - -@app.get("/mcp/resources") -async def get_mcp_resources(): - """Get available MCP resources""" - try: - from .mcp_integration import FuzeAgentMCPServer - - mcp_server = FuzeAgentMCPServer() - resources = [ - { - "uri": resource.uri, - "name": resource.name, - "description": resource.description, - "mime_type": resource.mime_type, - } - for resource in mcp_server.resources - ] - - return {"resources": resources} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get MCP resources: {str(e)}" - ) - - -@app.get("/mcp/resource") -async def get_mcp_resource(uri: str): - """Get an MCP resource by URI""" - try: - from .mcp_integration import FuzeAgentMCPServer - - if not uri: - raise HTTPException(status_code=400, detail="uri parameter is required") - - mcp_server = FuzeAgentMCPServer() - resource = await mcp_server.handle_resource_request(uri) - - return resource - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get MCP resource: {str(e)}" - ) - - -@app.get("/tasks/{task_id}/mcp-context") -async def get_task_mcp_context(task_id: str): - """Get MCP context for a task""" - try: - from .mcp_integration import FuzeAgentMCPServer, MCPClaudeIntegration - - execution = app.state.task_execution_engine.active_executions.get(task_id) - if not execution: - raise HTTPException(status_code=404, detail="Task not found or not active") - - mcp_server = FuzeAgentMCPServer() - mcp_integration = MCPClaudeIntegration(mcp_server) - - session_id = execution.claude_session_id or f"session-{task_id}" - context = await mcp_integration.get_session_context( - session_id=session_id, agent_id=execution.agent_id, task_id=task_id - ) - - return context - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get MCP context: {str(e)}" - ) - - -@app.post( - "/agents/{agent_id}/mcp-setup", - tags=["mcp-integration"], - summary="Setup Agent MCP Integration", - description="Configure MCP integration for an AI agent", -) -async def setup_agent_mcp( - agent_id: str = Path(..., description="Agent ID"), - setup_data: AgentMCPSetupRequest = Body(...), -): - """Set up MCP integration for an agent""" - try: - from .mcp_integration import FuzeAgentMCPServer, MCPClaudeIntegration - - task_id = setup_data.get("task_id") - session_id = setup_data.get("session_id") - - if not task_id: - raise HTTPException(status_code=400, detail="task_id is required") - - mcp_server = FuzeAgentMCPServer() - mcp_integration = MCPClaudeIntegration(mcp_server) - - # Set up MCP for Claude session - mcp_config = await mcp_integration.setup_claude_session_mcp( - session_id=session_id or f"session-{task_id}", - agent_id=agent_id, - task_id=task_id, - ) - - return { - "agent_id": agent_id, - "task_id": task_id, - "mcp_config": mcp_config, - "status": "mcp_configured", - } - - except Exception as e: - raise HTTPException(status_code=500, detail=f"Failed to setup MCP: {str(e)}") - - -# Multi-Agent Coordination Endpoints -@app.post( - "/tasks/{task_id}/coordinate", - tags=["multi-agent-coordination"], - summary="Initiate Multi-Agent Coordination", - description="Initiate multi-agent coordination for complex tasks", - response_model=CoordinationResponse, -) -async def initiate_task_coordination( - task_id: str = Path(..., description="Task ID to coordinate"), - coordination_request: CoordinationRequest = Body(...), -): - """Initiate multi-agent coordination for a complex task""" - try: - from .multi_agent_coordinator import CoordinationMode - - coordination_mode = coordination_request.get( - "coordination_mode", "collaborative" - ) - required_agents = coordination_request.get("required_agents") - required_skills = coordination_request.get("required_skills") - - # Validate coordination mode - try: - coord_mode = CoordinationMode(coordination_mode) - except ValueError: - raise HTTPException( - status_code=400, - detail=f"Invalid coordination mode: {coordination_mode}", - ) - - # Get multi-agent coordinator - coordinator = getattr( - app.state.task_execution_engine, "multi_agent_coordinator", None - ) - if not coordinator: - raise HTTPException( - status_code=503, detail="Multi-agent coordination not available" - ) - - # Initiate coordination - session_id = await coordinator.initiate_coordination( - task_id=task_id, - coordination_mode=coord_mode, - required_agents=required_agents, - required_skills=required_skills, - ) - - if session_id: - return { - "task_id": task_id, - "coordination_session_id": session_id, - "status": "coordination_initiated", - "coordination_mode": coordination_mode, - } - else: - return { - "task_id": task_id, - "status": "coordination_not_needed", - "message": "Task does not require multi-agent coordination", - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to initiate coordination: {str(e)}" - ) - - -@app.get("/coordination/{session_id}") -async def get_coordination_status(session_id: str): - """Get status of a coordination session""" - try: - coordinator = getattr( - app.state.task_execution_engine, "multi_agent_coordinator", None - ) - if not coordinator: - raise HTTPException( - status_code=503, detail="Multi-agent coordination not available" - ) - - status = await coordinator.get_coordination_status(session_id) - - if status: - return status - else: - raise HTTPException( - status_code=404, detail="Coordination session not found" - ) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get coordination status: {str(e)}" - ) - - -@app.post("/coordination/{session_id}/cancel") -async def cancel_coordination(session_id: str): - """Cancel a coordination session""" - try: - coordinator = getattr( - app.state.task_execution_engine, "multi_agent_coordinator", None - ) - if not coordinator: - raise HTTPException( - status_code=503, detail="Multi-agent coordination not available" - ) - - success = await coordinator.cancel_coordination(session_id) - - if success: - return {"coordination_session_id": session_id, "status": "cancelled"} - else: - raise HTTPException( - status_code=404, detail="Coordination session not found" - ) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to cancel coordination: {str(e)}" - ) - - -@app.post("/agents/{from_agent_id}/communicate/{to_agent_id}") -async def send_agent_communication( - from_agent_id: str, to_agent_id: str, communication_data: dict -): - """Send communication between agents""" - try: - message_type = communication_data.get("message_type", "notification") - content = communication_data.get("content", "") - metadata = communication_data.get("metadata", {}) - - if not content: - raise HTTPException(status_code=400, detail="Content cannot be empty") - - coordinator = getattr( - app.state.task_execution_engine, "multi_agent_coordinator", None - ) - if not coordinator: - raise HTTPException( - status_code=503, detail="Multi-agent coordination not available" - ) - - communication_id = await coordinator.send_agent_communication( - from_agent_id=from_agent_id, - to_agent_id=to_agent_id, - message_type=message_type, - content=content, - metadata=metadata, - ) - - return { - "communication_id": communication_id, - "from_agent_id": from_agent_id, - "to_agent_id": to_agent_id, - "status": "sent", - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to send agent communication: {str(e)}" - ) - - -@app.get("/coordination/active") -async def get_active_coordinations(): - """Get all active coordination sessions""" - try: - coordinator = getattr( - app.state.task_execution_engine, "multi_agent_coordinator", None - ) - if not coordinator: - raise HTTPException( - status_code=503, detail="Multi-agent coordination not available" - ) - - active_sessions = [] - for session_id in coordinator.active_sessions.keys(): - status = await coordinator.get_coordination_status(session_id) - if status: - active_sessions.append(status) - - return {"active_coordinations": active_sessions, "count": len(active_sessions)} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get active coordinations: {str(e)}" - ) - - -# WebSocket for coordination updates -@app.websocket("/ws/coordination/{session_id}") -async def coordination_websocket_endpoint(websocket: WebSocket, session_id: str): - """WebSocket endpoint for real-time coordination updates""" - await websocket.accept() - try: - coordinator = getattr( - app.state.task_execution_engine, "multi_agent_coordinator", None - ) - if not coordinator: - await websocket.send_json( - {"type": "error", "message": "Multi-agent coordination not available"} - ) - await websocket.close() - return - - # Monitor coordination session - while True: - try: - status = await coordinator.get_coordination_status(session_id) - if status: - await websocket.send_json( - { - "type": "coordination_update", - "session_id": session_id, - "data": status, - "timestamp": datetime.now().isoformat(), - } - ) - - # If coordination is completed or failed, send final update - if status.get("status") in ["completed", "failed", "cancelled"]: - await websocket.send_json( - { - "type": "coordination_finished", - "session_id": session_id, - "final_status": status.get("status"), - "timestamp": datetime.now().isoformat(), - } - ) - break - else: - await websocket.send_json( - { - "type": "error", - "message": f"Coordination session {session_id} not found", - } - ) - break - - await asyncio.sleep(3) # Update every 3 seconds - - except Exception as e: - await websocket.send_json( - { - "type": "error", - "message": f"Error monitoring coordination: {str(e)}", - } - ) - - except Exception as e: - print(f"Coordination WebSocket error for {session_id}: {e}") - finally: - await websocket.close() - - -# --------------------------------------------------------------------------- -# Agent relay WebSocket (Track 4) -# --------------------------------------------------------------------------- -@app.websocket("/agent-relay/{agent_id}") -async def agent_relay_endpoint(websocket: WebSocket, agent_id: str): - """ - Agent pods connect here to stream their session output. - Dashboard clients connect here to watch a specific agent's session. - Both use the same endpoint — first JSON message determines role: - {"role": "agent"} -> agent pod streaming output - {"role": "subscriber"} -> human dashboard watcher (default) - """ - await websocket.accept() - role = None - try: - init_msg = await websocket.receive_json() - role = init_msg.get("role", "subscriber") - - if role == "agent": - # Stream from agent pod to all subscribers - async for data in websocket.iter_json(): - msg = {"agentId": agent_id, **data} - dead = [] - for sub in list(agent_relay_subscribers[agent_id]): - try: - await sub.send_json(msg) - except Exception: - dead.append(sub) - for d in dead: - agent_relay_subscribers[agent_id].remove(d) - else: - # Human dashboard subscriber — wait for messages from agent - agent_relay_subscribers[agent_id].append(websocket) - await websocket.receive_text() # keep alive until disconnect - except WebSocketDisconnect: - pass - except Exception as e: - logger.warning(f"agent-relay {agent_id}: {e}") - finally: - subs = agent_relay_subscribers.get(agent_id, []) - if role != "agent" and websocket in subs: - subs.remove(websocket) - - -# Model Configuration and API Key Management Endpoints -@app.post( - "/organizations/{organization_id}/providers/{provider}/credentials", - tags=["model-configuration"], - summary="Store Provider API Credentials", - description="Store encrypted API credentials for a model provider", -) -async def store_provider_credentials( - organization_id: str = Path(..., description="Organization ID"), - provider: str = Path(..., description="Provider name"), - credentials: ProviderCredentialsRequest = Body(...), -): - """Store encrypted API credentials for a model provider at organization level""" - try: - from .model_configuration import ModelProvider, model_config_manager - - # Validate provider - try: - provider_enum = ModelProvider(provider) - except ValueError: - raise HTTPException( - status_code=400, detail=f"Unsupported provider: {provider}" - ) - - success = await model_config_manager.store_provider_credentials( - organization_id=organization_id, - provider=provider_enum, - api_key=credentials.api_key, - endpoint_url=credentials.endpoint_url, - additional_config=credentials.additional_config, - ) - - if success: - return { - "organization_id": organization_id, - "provider": provider, - "status": "credentials_stored", - "message": "API credentials stored successfully", - } - else: - raise HTTPException(status_code=500, detail="Failed to store credentials") - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to store provider credentials: {str(e)}" - ) - - -@app.get( - "/organizations/{organization_id}/models", - tags=["model-configuration"], - summary="Get Available Models", - description="Get available AI models for an organization", -) -async def get_available_models( - organization_id: str = Path(..., description="Organization ID"), - provider: Optional[str] = Query(None, description="Filter by provider"), - capabilities: Optional[str] = Query( - None, description="Filter by capabilities (comma-separated)" - ), -): - """Get available AI models with provider credential validation""" - try: - from .model_configuration import ( - ModelCapability, - ModelProvider, - model_config_manager, - ) - - provider_filter = None - if provider: - try: - provider_filter = ModelProvider(provider) - except ValueError: - raise HTTPException( - status_code=400, detail=f"Invalid provider: {provider}" - ) - - capabilities_filter = None - if capabilities: - try: - capabilities_filter = [ - ModelCapability(cap.strip()) for cap in capabilities.split(",") - ] - except ValueError as e: - raise HTTPException( - status_code=400, detail=f"Invalid capability: {str(e)}" - ) - - models = await model_config_manager.get_available_models( - organization_id=organization_id, - provider=provider_filter, - capabilities=capabilities_filter, - ) - - return { - "organization_id": organization_id, - "models": models, - "count": len(models), - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get available models: {str(e)}" - ) - - -@app.post( - "/agents/{agent_id}/model-configuration", - tags=["model-configuration"], - summary="Configure Agent Model Settings", - description="Configure model settings and preferences for an agent", -) -async def configure_agent_model( - agent_id: str = Path(..., description="Agent ID"), - config: AgentModelConfigRequest = Body(...), -): - """Configure model settings for an AI agent""" - try: - from .model_configuration import AgentModelConfig, model_config_manager - - agent_config = AgentModelConfig( - agent_id=agent_id, - primary_model=config.primary_model, - fallback_models=config.fallback_models, - temperature=config.temperature, - max_tokens=config.max_tokens, - top_p=config.top_p, - frequency_penalty=config.frequency_penalty, - presence_penalty=config.presence_penalty, - custom_instructions=config.custom_instructions, - use_function_calling=config.use_function_calling, - streaming_enabled=config.streaming_enabled, - cost_limit_per_task=config.cost_limit_per_task, - ) - - success = await model_config_manager.configure_agent_model( - agent_id, agent_config - ) - - if success: - return { - "agent_id": agent_id, - "status": "configured", - "primary_model": config.primary_model, - "fallback_models": config.fallback_models, - } - else: - raise HTTPException( - status_code=500, detail="Failed to configure agent model" - ) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to configure agent model: {str(e)}" - ) - - -@app.get( - "/agents/{agent_id}/model-configuration", - tags=["model-configuration"], - summary="Get Agent Model Configuration", - description="Get current model configuration for an agent", -) -async def get_agent_model_configuration( - agent_id: str = Path(..., description="Agent ID") -): - """Get model configuration for an AI agent""" - try: - from .model_configuration import model_config_manager - - config = await model_config_manager.get_agent_model_config(agent_id) - - if config: - return { - "agent_id": agent_id, - "configuration": { - "primary_model": config.primary_model, - "fallback_models": config.fallback_models, - "temperature": config.temperature, - "max_tokens": config.max_tokens, - "top_p": config.top_p, - "frequency_penalty": config.frequency_penalty, - "presence_penalty": config.presence_penalty, - "custom_instructions": config.custom_instructions, - "use_function_calling": config.use_function_calling, - "streaming_enabled": config.streaming_enabled, - "cost_limit_per_task": config.cost_limit_per_task, - "created_at": config.created_at.isoformat(), - "updated_at": config.updated_at.isoformat(), - }, - } - else: - raise HTTPException( - status_code=404, detail="Agent model configuration not found" - ) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get agent model configuration: {str(e)}" - ) - - -@app.post( - "/agents/{agent_id}/tasks/cost-estimate", - tags=["model-configuration"], - summary="Estimate Task Cost", - description="Estimate the cost of executing a task with the agent's model configuration", -) -async def estimate_task_cost( - agent_id: str = Path(..., description="Agent ID"), - request: TaskCostEstimateRequest = Body(...), -): - """Estimate cost for task execution based on agent's model configuration""" - try: - from .model_configuration import model_config_manager - - estimate = await model_config_manager.estimate_task_cost( - agent_id=agent_id, - task_description=request.task_description, - estimated_complexity=request.estimated_complexity, - ) - - return estimate - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to estimate task cost: {str(e)}" - ) - - -@app.get( - "/organizations/{organization_id}/model-usage", - tags=["model-configuration"], - summary="Get Model Usage Statistics", - description="Get model usage statistics and costs for an organization", -) -async def get_organization_model_usage( - organization_id: str = Path(..., description="Organization ID"), - days: int = Query(30, ge=1, le=365, description="Number of days to analyze"), -): - """Get model usage statistics and costs for an organization""" - try: - from .model_configuration import model_config_manager - - usage = await model_config_manager.get_organization_model_usage( - organization_id=organization_id, days=days - ) - - return usage - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get model usage: {str(e)}" - ) - - -@app.get( - "/agents/{agent_id}/model-recommendations", - tags=["model-configuration"], - summary="Get Model Recommendations", - description="Get model recommendations for an agent based on task capabilities", -) -async def get_model_recommendations( - agent_id: str = Path(..., description="Agent ID"), - capabilities: str = Query( - ..., description="Required capabilities (comma-separated)" - ), - cost_limit: Optional[float] = Query( - None, ge=0.0, description="Maximum cost limit in USD" - ), -): - """Get model recommendations based on task capabilities and cost constraints""" - try: - from .model_configuration import ModelCapability, model_config_manager - - # Parse capabilities - try: - capability_list = [ - ModelCapability(cap.strip()) for cap in capabilities.split(",") - ] - except ValueError as e: - raise HTTPException(status_code=400, detail=f"Invalid capability: {str(e)}") - - recommended_model = await model_config_manager.get_model_for_task( - agent_id=agent_id, task_capabilities=capability_list, cost_limit=cost_limit - ) - - if recommended_model: - return { - "agent_id": agent_id, - "recommended_model": recommended_model, - "capabilities": capabilities, - "cost_limit": cost_limit, - } - else: - return { - "agent_id": agent_id, - "recommended_model": None, - "message": "No suitable model found for the specified requirements", - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get model recommendations: {str(e)}" - ) - - -# Knowledge Management and Notification Endpoints - - -@app.get( - "/knowledge/notifications/{recipient_type}/{recipient_id}", - tags=["knowledge-management"], - summary="Get Knowledge Notifications", - description="Get notifications about knowledge updates, conflicts, and opportunities", -) -async def get_knowledge_notifications( - recipient_type: str = Path( - ..., description="Recipient type (agent, team, organization)" - ), - recipient_id: str = Path(..., description="Recipient ID"), - limit: int = Query(20, ge=1, le=100, description="Maximum notifications to return"), - status_filter: Optional[str] = Query( - None, description="Filter by status (unread, read, acknowledged)" - ), - notification_type_filter: Optional[str] = Query( - None, description="Filter by type (comma-separated)" - ), -): - """Get knowledge notifications for a recipient""" - try: - from .knowledge_notification_service import ( - KnowledgeNotificationService, - NotificationStatus, - NotificationType, - ) - - # Initialize notification service if not already done - if not hasattr(app.state, "notification_service"): - database_url = os.getenv( - "DATABASE_URL", - "postgresql://postgres:password@postgres:5432/ai_context", - ) - app.state.notification_service = KnowledgeNotificationService(database_url) - await app.state.notification_service.initialize() - - # Parse filters - status_filters = None - if status_filter: - try: - status_filters = [ - NotificationStatus(s.strip()) for s in status_filter.split(",") - ] - except ValueError as e: - raise HTTPException( - status_code=400, detail=f"Invalid status filter: {str(e)}" - ) - - type_filters = None - if notification_type_filter: - try: - type_filters = [ - NotificationType(t.strip()) - for t in notification_type_filter.split(",") - ] - except ValueError as e: - raise HTTPException( - status_code=400, - detail=f"Invalid notification type filter: {str(e)}", - ) - - notifications = ( - await app.state.notification_service.get_notifications_for_recipient( - recipient_type=recipient_type, - recipient_id=recipient_id, - limit=limit, - status_filter=status_filters, - notification_type_filter=type_filters, - ) - ) - - return { - "recipient_type": recipient_type, - "recipient_id": recipient_id, - "notifications": [ - { - "id": n.id, - "notification_type": n.notification_type.value, - "title": n.title, - "message": n.message, - "knowledge_id": n.knowledge_id, - "knowledge_type": n.knowledge_type, - "priority": n.priority.value, - "requires_action": n.requires_action, - "status": n.status.value, - "suggested_actions": n.suggested_actions, - "metadata": n.metadata, - "created_at": n.created_at.isoformat(), - "expires_at": n.expires_at.isoformat() if n.expires_at else None, - } - for n in notifications - ], - "count": len(notifications), - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get knowledge notifications: {str(e)}" - ) - - -@app.put( - "/knowledge/notifications/{notification_id}/status", - tags=["knowledge-management"], - summary="Update Notification Status", - description="Mark notification as read, acknowledged, or acted upon", -) -async def update_notification_status( - notification_id: str = Path(..., description="Notification ID"), - status: str = Body(..., description="New notification status"), - action_taken: Optional[Dict[str, Any]] = Body( - None, description="Optional action taken metadata" - ), -): - """Update notification status and optional action taken""" - try: - from .knowledge_notification_service import NotificationStatus - - # Validate status - try: - notification_status = NotificationStatus(status) - except ValueError: - raise HTTPException(status_code=400, detail=f"Invalid status: {status}") - - success = await app.state.notification_service.mark_notification_status( - notification_id=notification_id, - status=notification_status, - action_taken=action_taken, - ) - - if success: - return { - "notification_id": notification_id, - "status": status, - "updated": True, - } - else: - raise HTTPException(status_code=404, detail="Notification not found") - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to update notification status: {str(e)}" - ) - - -@app.get( - "/knowledge/notifications/statistics", - tags=["knowledge-management"], - summary="Get Notification Statistics", - description="Get comprehensive notification statistics and analytics", -) -async def get_notification_statistics( - organization_id: Optional[str] = Query( - None, description="Filter by organization ID" - ), - days_back: int = Query(30, ge=1, le=365, description="Days of history to analyze"), -): - """Get notification statistics and analytics""" - try: - stats = await app.state.notification_service.get_notification_statistics( - organization_id=organization_id, days_back=days_back - ) - - return stats - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get notification statistics: {str(e)}" - ) - - -@app.post( - "/knowledge/organizations/{organization_id}/add", - tags=["knowledge-management"], - summary="Add Organizational Knowledge", - description="Add knowledge to organization-level knowledge base", -) -async def add_organizational_knowledge( - organization_id: str = Path(..., description="Organization ID"), - title: str = Body(..., description="Knowledge title"), - content: str = Body(..., description="Knowledge content"), - content_type: str = Body("documentation", description="Content type"), - knowledge_category: str = Body("development", description="Knowledge category"), - source_agent_id: Optional[str] = Body(None, description="Source agent ID"), - source_team_id: Optional[str] = Body(None, description="Source team ID"), - tags: List[str] = Body(default_factory=list, description="Knowledge tags"), - metadata: Dict[str, Any] = Body( - default_factory=dict, description="Additional metadata" - ), -): - """Add knowledge to organization-level knowledge base""" - try: - from .organization_rag_manager import ( - ContentType, - KnowledgeCategory, - OrganizationRAGManager, - SourceType, - ) - - # Initialize services if not already done - if not hasattr(app.state, "org_rag_manager"): - database_url = os.getenv( - "DATABASE_URL", - "postgresql://postgres:password@postgres:5432/ai_context", - ) - app.state.org_rag_manager = OrganizationRAGManager(database_url) - await app.state.org_rag_manager.initialize() - - # Validate enums - try: - content_type_enum = ContentType(content_type) - category_enum = KnowledgeCategory(knowledge_category) - except ValueError as e: - raise HTTPException(status_code=400, detail=f"Invalid enum value: {str(e)}") - - knowledge_id = await app.state.org_rag_manager.add_knowledge( - organization_id=organization_id, - title=title, - content=content, - content_type=content_type_enum, - knowledge_category=category_enum, - source_type=SourceType.MANUAL_INPUT, - source_agent_id=source_agent_id, - source_team_id=source_team_id, - tags=tags, - metadata=metadata, - ) - - return { - "knowledge_id": knowledge_id, - "organization_id": organization_id, - "title": title, - "status": "added", - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to add organizational knowledge: {str(e)}" - ) - - -@app.get( - "/knowledge/organizations/{organization_id}/search", - tags=["knowledge-management"], - summary="Search Organizational Knowledge", - description="Search organization-level knowledge base", -) -async def search_organizational_knowledge( - organization_id: str = Path(..., description="Organization ID"), - query: str = Query(..., description="Search query"), - limit: int = Query(10, ge=1, le=50, description="Maximum results to return"), - min_similarity: float = Query( - 0.3, ge=0.0, le=1.0, description="Minimum similarity threshold" - ), - categories: Optional[str] = Query( - None, description="Filter by categories (comma-separated)" - ), -): - """Search organization-level knowledge base""" - try: - from .organization_rag_manager import KnowledgeCategory - - # Parse categories - category_filters = None - if categories: - try: - category_filters = [ - KnowledgeCategory(cat.strip()) for cat in categories.split(",") - ] - except ValueError as e: - raise HTTPException( - status_code=400, detail=f"Invalid category: {str(e)}" - ) - - search_results = await app.state.org_rag_manager.search_knowledge( - organization_id=organization_id, - query=query, - limit=limit, - min_similarity=min_similarity, - categories=category_filters, - ) - - results = [] - for result in search_results: - results.append( - { - "knowledge_id": result.knowledge.id, - "title": result.knowledge.title, - "content_preview": ( - result.knowledge.content[:200] + "..." - if len(result.knowledge.content) > 200 - else result.knowledge.content - ), - "category": result.knowledge.knowledge_category.value, - "content_type": result.knowledge.content_type.value, - "similarity_score": result.similarity_score, - "combined_score": result.combined_score, - "quality_score": result.knowledge.quality_score, - "usage_count": result.knowledge.usage_count, - "created_at": result.knowledge.created_at.isoformat(), - "tags": result.knowledge.tags, - "metadata": result.knowledge.metadata, - } - ) - - return { - "organization_id": organization_id, - "query": query, - "results": results, - "count": len(results), - } - - except Exception as e: - raise HTTPException( - status_code=500, - detail=f"Failed to search organizational knowledge: {str(e)}", - ) - - -@app.get( - "/knowledge/context-enhancement/{agent_id}", - tags=["knowledge-management"], - summary="Get Enhanced Context for Agent", - description="Get enhanced context with relevant organizational knowledge for task execution", -) -async def get_enhanced_context_for_agent( - agent_id: str = Path(..., description="Agent ID"), - task_description: str = Query( - ..., description="Task description for context enhancement" - ), - task_type: Optional[str] = Query(None, description="Task type"), - technologies: Optional[str] = Query( - None, description="Technologies involved (comma-separated)" - ), -): - """Get enhanced context with relevant knowledge for agent task execution""" - try: - from .context_enhancement_service import ContextEnhancementService - - # Initialize context enhancement service if needed - if not hasattr(app.state, "context_enhancement_service"): - database_url = os.getenv( - "DATABASE_URL", - "postgresql://postgres:password@postgres:5432/ai_context", - ) - # These would be initialized in the lifespan - if hasattr(app.state, "org_rag_manager") and hasattr( - app.state, "team_knowledge_manager" - ): - app.state.context_enhancement_service = ContextEnhancementService( - database_url=database_url, - org_rag_manager=app.state.org_rag_manager, - team_knowledge_manager=app.state.team_knowledge_manager, - ) - await app.state.context_enhancement_service.initialize() - else: - raise HTTPException( - status_code=503, - detail="Knowledge management services not initialized", - ) - - # Build task data - task_data = { - "description": task_description, - "task_type": task_type, - "technologies": technologies.split(",") if technologies else [], - } - - enhanced_context = ( - await app.state.context_enhancement_service.enhance_agent_context( - agent_id=agent_id, task_data=task_data - ) - ) - - return { - "agent_id": agent_id, - "task_description": task_description, - "enhanced_context": { - "organizational_knowledge_count": len( - enhanced_context.organizational_knowledge - ), - "team_knowledge_count": len(enhanced_context.team_knowledge), - "similar_task_insights_count": len( - enhanced_context.similar_task_insights - ), - "success_patterns": enhanced_context.success_patterns, - "common_pitfalls": enhanced_context.common_pitfalls, - "recommended_approaches": enhanced_context.recommended_approaches, - "context_summary": enhanced_context.context_summary, - "enhancement_metadata": enhanced_context.enhancement_metadata, - }, - "organizational_knowledge": [ - { - "knowledge_id": item.knowledge_id, - "title": item.title, - "category": item.category, - "relevance_score": item.relevance_score, - "confidence_score": item.confidence_score, - "content_preview": ( - item.content[:200] + "..." - if len(item.content) > 200 - else item.content - ), - } - for item in enhanced_context.organizational_knowledge - ], - "team_knowledge": [ - { - "knowledge_id": item.knowledge_id, - "title": item.title, - "category": item.category, - "relevance_score": item.relevance_score, - "confidence_score": item.confidence_score, - "content_preview": ( - item.content[:200] + "..." - if len(item.content) > 200 - else item.content - ), - } - for item in enhanced_context.team_knowledge - ], - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get enhanced context: {str(e)}" - ) - - -@app.get( - "/knowledge/analytics/organizations/{organization_id}/insights", - tags=["knowledge-management"], - summary="Get Organizational Knowledge Insights", - description="Get comprehensive analytics and insights about organizational knowledge", -) -async def get_organizational_knowledge_insights( - organization_id: str = Path(..., description="Organization ID"), - analysis_period_days: int = Query( - 30, ge=7, le=365, description="Analysis period in days" - ), -): - """Get comprehensive organizational knowledge insights and analytics""" - try: - insights = ( - await app.state.knowledge_analytics_service.get_organizational_insights( - organization_id=organization_id, - analysis_period_days=analysis_period_days, - ) - ) - - return { - "organization_id": organization_id, - "analysis_period_days": analysis_period_days, - "insights": { - "total_knowledge_items": insights.total_knowledge_items, - "knowledge_growth_rate": insights.knowledge_growth_rate, - "knowledge_utilization_rate": insights.knowledge_utilization_rate, - "knowledge_freshness_score": insights.knowledge_freshness_score, - "cross_team_sharing_rate": insights.cross_team_sharing_rate, - "propagation_efficiency": insights.propagation_efficiency, - "top_performing_categories": insights.top_performing_categories, - "knowledge_gaps": insights.knowledge_gaps, - "agent_knowledge_engagement": insights.agent_knowledge_engagement, - "team_knowledge_contribution": insights.team_knowledge_contribution, - "recommendations": insights.recommendations, - }, - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get organizational insights: {str(e)}" - ) - - -@app.get( - "/knowledge/analytics/organizations/{organization_id}/effectiveness", - tags=["knowledge-management"], - summary="Analyze Knowledge Effectiveness", - description="Analyze effectiveness and performance of knowledge items", -) -async def analyze_knowledge_effectiveness( - organization_id: str = Path(..., description="Organization ID"), - knowledge_category: Optional[str] = Query( - None, description="Filter by knowledge category" - ), - min_usage_count: int = Query( - 3, ge=1, description="Minimum usage count for analysis" - ), -): - """Analyze effectiveness of knowledge items in the organization""" - try: - effectiveness_metrics = ( - await app.state.knowledge_analytics_service.analyze_knowledge_effectiveness( - organization_id=organization_id, - knowledge_category=knowledge_category, - min_usage_count=min_usage_count, - ) - ) - - results = [] - for metric in effectiveness_metrics: - results.append( - { - "knowledge_id": metric.knowledge_id, - "title": metric.title, - "category": metric.category, - "usage_count": metric.usage_count, - "success_correlation": metric.success_correlation, - "average_relevance": metric.average_relevance, - "agent_adoption_rate": metric.agent_adoption_rate, - "team_adoption_rate": metric.team_adoption_rate, - "quality_score": metric.quality_score, - "recency_score": metric.recency_score, - "overall_effectiveness": metric.overall_effectiveness, - "trend_direction": metric.trend_direction, - "optimization_suggestions": metric.optimization_suggestions, - } - ) - - return { - "organization_id": organization_id, - "effectiveness_analysis": results, - "total_analyzed": len(results), - "summary": { - "avg_effectiveness": sum(r["overall_effectiveness"] for r in results) - / max(len(results), 1), - "top_performers": sorted( - results, key=lambda x: x["overall_effectiveness"], reverse=True - )[:5], - "needs_attention": [ - r for r in results if r["overall_effectiveness"] < 0.5 - ], - }, - } - - except Exception as e: - raise HTTPException( - status_code=500, - detail=f"Failed to analyze knowledge effectiveness: {str(e)}", - ) - - -@app.get( - "/knowledge/analytics/agents/{agent_id}/profile", - tags=["knowledge-management"], - summary="Get Agent Knowledge Profile", - description="Get detailed knowledge profile and analytics for an agent", -) -async def get_agent_knowledge_profile( - agent_id: str = Path(..., description="Agent ID"), - analysis_period_days: int = Query( - 60, ge=7, le=365, description="Analysis period in days" - ), -): - """Get detailed knowledge profile for an agent""" - try: - profile = ( - await app.state.knowledge_analytics_service.get_agent_knowledge_profile( - agent_id=agent_id, analysis_period_days=analysis_period_days - ) - ) - - if not profile: - raise HTTPException( - status_code=404, detail="Agent not found or no knowledge data available" - ) - - return { - "agent_id": agent_id, - "analysis_period_days": analysis_period_days, - "profile": { - "agent_name": profile.agent_name, - "team_id": profile.team_id, - "knowledge_consumption_rate": profile.knowledge_consumption_rate, - "knowledge_creation_rate": profile.knowledge_creation_rate, - "expertise_areas": profile.expertise_areas, - "knowledge_application_success": profile.knowledge_application_success, - "learning_velocity": profile.learning_velocity, - "knowledge_sharing_activity": profile.knowledge_sharing_activity, - "preferred_knowledge_types": profile.preferred_knowledge_types, - "knowledge_gaps": profile.knowledge_gaps, - "optimization_recommendations": profile.optimization_recommendations, - }, - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get agent knowledge profile: {str(e)}" - ) - - -@app.get( - "/knowledge/analytics/organizations/{organization_id}/optimization", - tags=["knowledge-management"], - summary="Get Knowledge Optimization Recommendations", - description="Get comprehensive recommendations for knowledge system optimization", -) -async def get_knowledge_optimization_recommendations( - organization_id: str = Path(..., description="Organization ID"), - focus_area: Optional[str] = Query( - None, - description="Focus area (utilization, quality, gaps, propagation, collaboration)", - ), -): - """Generate comprehensive knowledge optimization recommendations""" - try: - recommendations = await app.state.knowledge_analytics_service.generate_knowledge_optimization_recommendations( - organization_id=organization_id, focus_area=focus_area - ) - - return { - "organization_id": organization_id, - "focus_area": focus_area, - "recommendations": recommendations, - "total_recommendations": len(recommendations), - } - - except Exception as e: - raise HTTPException( - status_code=500, - detail=f"Failed to get optimization recommendations: {str(e)}", - ) - - -@app.get( - "/knowledge/analytics/organizations/{organization_id}/trends", - tags=["knowledge-management"], - summary="Get Knowledge Trends Analysis", - description="Analyze knowledge trends and patterns over time", -) -async def get_knowledge_trends_analysis( - organization_id: str = Path(..., description="Organization ID"), - trend_period_days: int = Query( - 90, ge=30, le=365, description="Trend analysis period in days" - ), -): - """Get comprehensive knowledge trends analysis""" - try: - trends = ( - await app.state.knowledge_analytics_service.get_knowledge_trends_analysis( - organization_id=organization_id, trend_period_days=trend_period_days - ) - ) - - return { - "organization_id": organization_id, - "trend_period_days": trend_period_days, - "trends": trends, - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get knowledge trends: {str(e)}" - ) - - -# Memory-Enhanced Agents Endpoints - - -@app.post( - "/agents/{agent_id}/deploy-memory", - tags=["memory-agents"], - summary="Deploy Memory-Enabled Agent", - description="Deploy an agent with persistent memory capabilities", -) -async def deploy_memory_enabled_agent( - agent_id: str = Path(..., description="Agent ID"), - template_id: str = Body(..., description="Agent template ID"), - task_id: Optional[str] = Body(None, description="Optional specific task ID"), - repository_settings: Optional[Dict[str, Any]] = Body( - None, description="Repository settings" - ), -): - """Deploy a memory-enabled autonomous agent container""" - try: - result = await app.state.agent_manager.deploy_memory_enabled_agent( - agent_id=agent_id, - template_id=template_id, - task_id=task_id, - repository_settings=repository_settings, - ) - - if result["success"]: - return result - else: - raise HTTPException(status_code=500, detail=result["error"]) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to deploy memory-enabled agent: {str(e)}" - ) - - -@app.get( - "/agents/{agent_id}/memory-status", - tags=["memory-agents"], - summary="Get Agent Memory Status", - description="Get agent memory status and expertise summary", -) -async def get_agent_memory_status(agent_id: str = Path(..., description="Agent ID")): - """Get agent memory status, expertise metrics, and insights""" - try: - status = await app.state.agent_manager.get_agent_memory_status(agent_id) - return status - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get agent memory status: {str(e)}" - ) - - -@app.post( - "/agents/{agent_id}/memory-tasks", - tags=["memory-agents"], - summary="Assign Task to Memory Agent", - description="Assign a task to a memory-enabled agent", -) -async def assign_task_to_memory_agent( - agent_id: str = Path(..., description="Agent ID"), - task_id: str = Body(..., description="Task ID"), - task_data: Dict[str, Any] = Body(..., description="Task data"), -): - """Assign a task to a memory-enabled agent for autonomous execution""" - try: - result = await app.state.agent_manager.assign_task_to_memory_agent( - agent_id=agent_id, task_id=task_id, task_data=task_data - ) - - if result["success"]: - return result - else: - raise HTTPException(status_code=400, detail=result["error"]) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to assign task to memory agent: {str(e)}" - ) - - -@app.delete( - "/agents/{agent_id}/memory", - tags=["memory-agents"], - summary="Stop Memory-Enabled Agent", - description="Stop a memory-enabled agent container", -) -async def stop_memory_enabled_agent(agent_id: str = Path(..., description="Agent ID")): - """Stop and clean up a memory-enabled agent container""" - try: - result = await app.state.agent_manager.stop_memory_enabled_agent(agent_id) - - if result["success"]: - return result - else: - raise HTTPException(status_code=400, detail=result["error"]) - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to stop memory-enabled agent: {str(e)}" - ) - - -@app.get( - "/system/expertise-dashboard", - tags=["memory-agents"], - summary="Get System Expertise Dashboard", - description="Get system-wide expertise and memory analytics", -) -async def get_system_expertise_dashboard(): - """Get comprehensive dashboard of system expertise and memory analytics""" - try: - dashboard = await app.state.agent_manager.get_system_expertise_dashboard() - return dashboard - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get expertise dashboard: {str(e)}" - ) - - -@app.get( - "/agents/{agent_id}/tasks/pending", - tags=["memory-agents"], - summary="Get Pending Tasks for Agent", - description="Get pending tasks for a memory-enabled agent", -) -async def get_pending_tasks_for_agent( - agent_id: str = Path(..., description="Agent ID"), - limit: int = Query( - 10, ge=1, le=50, description="Maximum number of tasks to return" - ), -): - """Get pending tasks that a memory-enabled agent can pick up""" - try: - async with get_db_connection() as conn: - tasks = await conn.fetch( - """ - SELECT id, title, description, type, complexity, language, - framework, requirements, created_at - FROM tasks - WHERE agent_id = $1 - AND status = 'pending' - AND assigned_to_memory_agent = true - ORDER BY created_at ASC - LIMIT $2 - """, - agent_id, - limit, - ) - - return { - "agent_id": agent_id, - "tasks": [dict(task) for task in tasks], - "count": len(tasks), - } - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to get pending tasks: {str(e)}" - ) - - -@app.put( - "/tasks/{task_id}/status", - tags=["memory-agents"], - summary="Update Task Status", - description="Update task status (used by memory-enabled agents)", -) -async def update_task_status( - task_id: str = Path(..., description="Task ID"), - status: str = Body(..., description="New task status"), - result: Optional[Dict[str, Any]] = Body(None, description="Task result data"), - updated_by: Optional[str] = Body(None, description="ID of agent updating the task"), - container_instance_id: Optional[str] = Body( - None, description="Container instance ID" - ), - updated_at: Optional[str] = Body(None, description="Update timestamp"), -): - """Update task status - used by memory-enabled agents to report progress""" - try: - async with get_db_connection() as conn: - await conn.execute( - """ - UPDATE tasks - SET status = $2, - result = COALESCE($3, result), - updated_by = COALESCE($4, updated_by), - updated_at = NOW() - WHERE id = $1 - """, - task_id, - status, - result, - updated_by, - ) - - # If task is completed, log it for expertise tracking - if status in ["completed", "failed"]: - # The agent's memory system will handle learning from the outcome - pass - - return {"task_id": task_id, "status": status, "updated": True} - - except Exception as e: - raise HTTPException( - status_code=500, detail=f"Failed to update task status: {str(e)}" - ) - - -@app.post( - "/agents/{agent_id}/register", - tags=["memory-agents"], - summary="Agent Registration", - description="Register agent capabilities and status with orchestrator", -) -async def register_agent_capabilities( - agent_id: str = Path(..., description="Agent ID"), - capabilities: Dict[str, Any] = Body( - ..., description="Agent capabilities and status" - ), -): - """Register or update agent capabilities - used by memory-enabled agents on startup""" - try: - # Update agent capabilities in database - async with get_db_connection() as conn: - await conn.execute( - """ - UPDATE agents - SET config = config || $2, - status = 'active', - updated_at = NOW() - WHERE id = $1 - """, - agent_id, - { - "capabilities": capabilities, - "last_registration": datetime.now().isoformat(), - }, - ) - - # Update in-memory tracking - if agent_id in app.state.agent_manager.memory_enabled_agents: - app.state.agent_manager.memory_enabled_agents[agent_id]["status"] = "active" - - return { - "agent_id": agent_id, - "agent_recognized": True, - "capabilities_accepted": True, - "status": "registered", - } - - except Exception as e: - return { - "agent_id": agent_id, - "agent_recognized": False, - "capabilities_accepted": False, - "error": str(e), - } - - -@app.post( - "/agents/{agent_id}/statistics", - tags=["memory-agents"], - summary="Agent Statistics Update", - description="Update agent performance and memory statistics", -) -async def update_agent_statistics( - agent_id: str = Path(..., description="Agent ID"), - stats: Dict[str, Any] = Body(..., description="Agent statistics"), -): - """Update agent statistics - used by memory-enabled agents for performance tracking""" - try: - # Store statistics for analytics - async with get_db_connection() as conn: - await conn.execute( - """ - UPDATE agents - SET config = config || $2, - updated_at = NOW() - WHERE id = $1 - """, - agent_id, - { - "latest_statistics": stats, - "statistics_updated_at": datetime.now().isoformat(), - }, - ) - - # Clear expertise cache to force refresh - await app.state.agent_manager.expertise_tracker.clear_cache(agent_id) - - return {"agent_id": agent_id, "statistics_updated": True} - - except Exception as e: - return {"agent_id": agent_id, "statistics_updated": False, "error": str(e)} - - -@app.post( - "/agents/{agent_id}/error", - tags=["memory-agents"], - summary="Agent Error Reporting", - description="Report agent errors for monitoring", -) -async def report_agent_error( - agent_id: str = Path(..., description="Agent ID"), - error_data: Dict[str, Any] = Body(..., description="Error information"), -): - """Report agent errors - used by memory-enabled agents for error tracking""" - try: - # Log error for monitoring - logger.error(f"Agent {agent_id} reported error: {error_data}") - - # Update agent status if it's a critical error - if error_data.get("critical", False): - async with get_db_connection() as conn: - await conn.execute( - """ - UPDATE agents - SET status = 'error', - config = config || $2, - updated_at = NOW() - WHERE id = $1 - """, - agent_id, - { - "last_error": error_data, - "error_reported_at": datetime.now().isoformat(), - }, - ) - - return {"agent_id": agent_id, "error_logged": True} - - except Exception as e: - logger.error(f"Failed to log agent error: {e}") - return {"agent_id": agent_id, "error_logged": False} - - -# ============================================================================ -# Goals Management API Endpoints -# ============================================================================ - - -@app.post( - "/organizations/{organization_id}/goals", - tags=["goals-management"], - summary="Create organizational goal", - description="Create a new goal for an organization with specified targets and deadlines", -) -async def create_goal( - organization_id: str = Path(..., description="Organization ID"), - goal_data: GoalCreateRequest = Body(..., description="Goal creation data"), - created_by: Optional[str] = Query( - None, description="ID of user/agent creating the goal" - ), -): - """Create a new organizational goal""" - try: - from .goals_management_service import GoalType - - goal_id = await app.state.goals_service.create_goal( - organization_id=organization_id, - title=goal_data.title, - description=goal_data.description, - goal_type=GoalType(goal_data.goal_type), - target_value=goal_data.target_value, - target_unit=goal_data.target_unit, - target_deadline=goal_data.target_deadline, - priority_level=goal_data.priority_level, - success_criteria=goal_data.success_criteria, - assigned_teams=goal_data.assigned_teams, - goal_owner_agent_id=goal_data.goal_owner_agent_id, - stakeholder_agents=goal_data.stakeholder_agents, - tags=goal_data.tags, - metadata=goal_data.metadata, - created_by=created_by, - ) - - return {"goal_id": goal_id, "status": "created"} - - except Exception as e: - logger.error(f"Error creating goal: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/organizations/{organization_id}/goals", - tags=["goals-management"], - summary="List organization goals", - description="Get all goals for an organization with optional filtering", -) -async def list_organization_goals( - organization_id: str = Path(..., description="Organization ID"), - status: Optional[List[str]] = Query(None, description="Filter by goal status"), - goal_type: Optional[List[str]] = Query(None, description="Filter by goal type"), - limit: int = Query( - 50, ge=1, le=100, description="Maximum number of goals to return" - ), -): - """List goals for an organization""" - try: - from .goals_management_service import GoalStatus, GoalType - - status_filter = [GoalStatus(s) for s in status] if status else None - type_filter = [GoalType(gt) for gt in goal_type] if goal_type else None - - goals = await app.state.goals_service.list_organization_goals( - organization_id=organization_id, - status_filter=status_filter, - goal_type_filter=type_filter, - limit=limit, - ) - - return { - "organization_id": organization_id, - "goals": [ - { - "id": goal.id, - "title": goal.title, - "description": goal.description, - "goal_type": goal.goal_type.value, - "status": goal.status.value, - "progress_percentage": float(goal.progress_percentage), - "target_value": ( - float(goal.target_value) if goal.target_value else None - ), - "target_unit": goal.target_unit, - "current_value": ( - float(goal.current_value) if goal.current_value else None - ), - "target_deadline": goal.target_deadline.isoformat(), - "priority_level": goal.priority_level, - "completion_confidence": float(goal.completion_confidence), - "created_at": goal.created_at.isoformat(), - "updated_at": goal.updated_at.isoformat(), - } - for goal in goals - ], - } - - except Exception as e: - logger.error(f"Error listing organization goals: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/goals/{goal_id}", - tags=["goals-management"], - summary="Get goal details", - description="Get detailed information about a specific goal", -) -async def get_goal(goal_id: str = Path(..., description="Goal ID")): - """Get goal details""" - try: - goal = await app.state.goals_service.get_goal(goal_id) - - if not goal: - raise HTTPException(status_code=404, detail="Goal not found") - - return { - "id": goal.id, - "organization_id": goal.organization_id, - "title": goal.title, - "description": goal.description, - "goal_type": goal.goal_type.value, - "status": goal.status.value, - "progress_percentage": float(goal.progress_percentage), - "target_value": float(goal.target_value) if goal.target_value else None, - "target_unit": goal.target_unit, - "current_value": float(goal.current_value) if goal.current_value else None, - "success_criteria": goal.success_criteria, - "start_date": goal.start_date.isoformat(), - "target_deadline": goal.target_deadline.isoformat(), - "actual_completion_date": ( - goal.actual_completion_date.isoformat() - if goal.actual_completion_date - else None - ), - "priority_level": goal.priority_level, - "completion_confidence": float(goal.completion_confidence), - "assigned_teams": goal.assigned_teams, - "goal_owner_agent_id": goal.goal_owner_agent_id, - "stakeholder_agents": goal.stakeholder_agents, - "tags": goal.tags, - "metadata": goal.metadata, - "created_by": goal.created_by, - "created_at": goal.created_at.isoformat(), - "updated_at": goal.updated_at.isoformat(), - } - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting goal {goal_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/goals/{goal_id}/overview", - tags=["goals-management"], - summary="Get goal overview", - description="Get comprehensive overview of goal with milestones, tasks, and progress", -) -async def get_goal_overview(goal_id: str = Path(..., description="Goal ID")): - """Get comprehensive goal overview""" - try: - overview = await app.state.goals_service.get_goal_overview(goal_id) - - if not overview: - raise HTTPException(status_code=404, detail="Goal not found") - - return overview - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting goal overview {goal_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.put( - "/goals/{goal_id}/progress", - tags=["goals-management"], - summary="Update goal progress", - description="Update progress for a specific goal", -) -async def update_goal_progress( - goal_id: str = Path(..., description="Goal ID"), - progress_data: GoalUpdateRequest = Body(..., description="Progress update data"), - recorded_by: Optional[str] = Query( - None, description="ID of user/agent recording progress" - ), -): - """Update goal progress""" - try: - success = await app.state.goals_service.update_goal_progress( - goal_id=goal_id, - progress_percentage=progress_data.progress_percentage, - current_value=progress_data.current_value, - completion_confidence=progress_data.completion_confidence, - progress_notes=progress_data.notes, - recorded_by=recorded_by, - ) - - if not success: - raise HTTPException( - status_code=404, detail="Goal not found or no changes made" - ) - - return {"goal_id": goal_id, "status": "updated"} - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error updating goal progress {goal_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/goals/{goal_id}/milestones", - tags=["goals-management"], - summary="Create milestone", - description="Create a new milestone for a goal", -) -async def create_milestone( - goal_id: str = Path(..., description="Goal ID"), - milestone_data: MilestoneCreateRequest = Body( - ..., description="Milestone creation data" - ), - created_by: Optional[str] = Query( - None, description="ID of user/agent creating milestone" - ), -): - """Create milestone for goal""" - try: - milestone_id = await app.state.goals_service.create_milestone( - goal_id=goal_id, - title=milestone_data.title, - description=milestone_data.description, - target_date=milestone_data.target_date, - milestone_type=milestone_data.milestone_type, - success_criteria=milestone_data.success_criteria, - deliverables=milestone_data.deliverables, - dependencies=milestone_data.dependencies, - assigned_teams=milestone_data.assigned_teams, - responsible_agent_id=milestone_data.responsible_agent_id, - priority_level=milestone_data.priority_level, - weight_in_goal=milestone_data.weight_in_goal, - created_by=created_by, - ) - - return {"milestone_id": milestone_id, "status": "created"} - - except Exception as e: - logger.error(f"Error creating milestone: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/milestones/{milestone_id}/tasks", - tags=["goals-management"], - summary="Create task from milestone", - description="Create a new task derived from a milestone", -) -async def create_task_from_milestone( - milestone_id: str = Path(..., description="Milestone ID"), - task_data: TaskFromMilestoneRequest = Body(..., description="Task creation data"), - created_by: Optional[str] = Query( - None, description="ID of user/agent creating task" - ), -): - """Create task from milestone""" - try: - task_id = await app.state.goals_service.create_task_from_milestone( - milestone_id=milestone_id, - title=task_data.title, - description=task_data.description, - task_type=task_data.task_type, - complexity_level=task_data.complexity_level, - estimated_hours=task_data.estimated_hours, - due_date=task_data.due_date, - assigned_team_id=task_data.assigned_team_id, - assigned_agent_id=task_data.assigned_agent_id, - priority=task_data.priority, - requirements=task_data.requirements, - acceptance_criteria=task_data.acceptance_criteria, - dependencies=task_data.dependencies, - created_by_agent_id=created_by, - ) - - return {"task_id": task_id, "status": "created"} - - except Exception as e: - logger.error(f"Error creating task from milestone: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/goals/{goal_id}/generate-execution-plan", - tags=["goals-management"], - summary="Generate execution plan", - description="Generate comprehensive milestone and task execution plan for a goal", -) -async def generate_execution_plan( - goal_id: str = Path(..., description="Goal ID"), - planning_context: Optional[Dict[str, Any]] = Body( - None, description="Additional planning context" - ), -): - """Generate execution plan with milestones and tasks""" - try: - execution_plan = ( - await app.state.milestone_task_engine.generate_goal_execution_plan( - goal_id=goal_id, planning_context=planning_context - ) - ) - - return execution_plan - - except Exception as e: - logger.error(f"Error generating execution plan for goal {goal_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/goals/{goal_id}/generate-monthly-milestones", - tags=["goals-management"], - summary="Generate monthly milestones", - description="Generate monthly milestone breakdown for a goal", -) -async def generate_monthly_milestones( - goal_id: str = Path(..., description="Goal ID"), - start_date: Optional[date] = Query(None, description="Start date for milestones"), - end_date: Optional[date] = Query(None, description="End date for milestones"), -): - """Generate monthly milestones for goal""" - try: - milestone_ids = ( - await app.state.milestone_task_engine.generate_monthly_milestones( - goal_id=goal_id, start_date=start_date, end_date=end_date - ) - ) - - return { - "goal_id": goal_id, - "milestone_ids": milestone_ids, - "count": len(milestone_ids), - "status": "generated", - } - - except Exception as e: - logger.error(f"Error generating monthly milestones for goal {goal_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/milestones/{milestone_id}/generate-weekly-tasks", - tags=["goals-management"], - summary="Generate weekly tasks", - description="Generate weekly task breakdown for a milestone", -) -async def generate_weekly_tasks( - milestone_id: str = Path(..., description="Milestone ID"), - focus_areas: Optional[List[str]] = Body( - None, description="Focus areas for task generation" - ), -): - """Generate weekly tasks for milestone""" - try: - task_ids = ( - await app.state.milestone_task_engine.generate_weekly_tasks_for_milestone( - milestone_id=milestone_id, focus_areas=focus_areas - ) - ) - - return { - "milestone_id": milestone_id, - "task_ids": task_ids, - "count": len(task_ids), - "status": "generated", - } - - except Exception as e: - logger.error(f"Error generating weekly tasks for milestone {milestone_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/goals/{goal_id}/generate-cross-functional-tasks", - tags=["goals-management"], - summary="Generate cross-functional tasks", - description="Generate tasks across different business functions for a goal", -) -async def generate_cross_functional_tasks( - goal_id: str = Path(..., description="Goal ID"), - target_functions: Optional[List[str]] = Body( - None, description="Target business functions" - ), -): - """Generate cross-functional tasks for goal""" - try: - functional_tasks = ( - await app.state.milestone_task_engine.generate_cross_functional_tasks( - goal_id=goal_id, target_functions=target_functions - ) - ) - - return { - "goal_id": goal_id, - "functional_tasks": functional_tasks, - "total_tasks": sum(len(tasks) for tasks in functional_tasks.values()), - "status": "generated", - } - - except Exception as e: - logger.error(f"Error generating cross-functional tasks for goal {goal_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/goals/{goal_id}/conversations", - tags=["goals-management"], - summary="Create goal conversation", - description="Create AI-powered conversation for goal planning and discussion", -) -async def create_goal_conversation( - goal_id: str = Path(..., description="Goal ID"), - conversation_data: GoalConversationCreateRequest = Body( - ..., description="Conversation creation data" - ), - created_by: Optional[str] = Query( - None, description="ID of user/agent creating conversation" - ), -): - """Create goal conversation""" - try: - from .goal_conversation_service import ConversationType - - conversation_id = ( - await app.state.goal_conversation_service.create_goal_conversation( - goal_id=goal_id, - conversation_type=ConversationType(conversation_data.conversation_type), - conversation_title=conversation_data.conversation_title, - initial_context=conversation_data.initial_context, - participants=conversation_data.participants, - created_by=created_by, - ) - ) - - return {"conversation_id": conversation_id, "status": "created"} - - except Exception as e: - logger.error(f"Error creating goal conversation: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/conversations/{conversation_id}", - tags=["goals-management"], - summary="Get goal conversation", - description="Get full conversation with messages, insights, and action items", -) -async def get_goal_conversation( - conversation_id: str = Path(..., description="Conversation ID") -): - """Get goal conversation""" - try: - conversation = await app.state.goal_conversation_service.get_conversation( - conversation_id - ) - - if not conversation: - raise HTTPException(status_code=404, detail="Conversation not found") - - return conversation - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting conversation {conversation_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/conversations/{conversation_id}/messages", - tags=["goals-management"], - summary="Add message to conversation", - description="Add a new message to a goal conversation", -) -async def add_message_to_conversation( - conversation_id: str = Path(..., description="Conversation ID"), - message_data: ConversationMessageRequest = Body(..., description="Message data"), - sender_id: Optional[str] = Query(None, description="ID of message sender"), -): - """Add message to conversation""" - try: - from .goal_conversation_service import MessageType - - message_id = ( - await app.state.goal_conversation_service.add_message_to_conversation( - conversation_id=conversation_id, - message_type=MessageType(message_data.message_type), - sender_id=sender_id, - sender_name=message_data.sender_name, - content=message_data.content, - metadata=message_data.metadata, - references=message_data.references, - ) - ) - - return {"message_id": message_id, "status": "added"} - - except Exception as e: - logger.error(f"Error adding message to conversation: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/conversations/{conversation_id}/generate-milestones", - tags=["goals-management"], - summary="Generate milestones from conversation", - description="Generate milestone recommendations based on conversation analysis", -) -async def generate_planning_milestones( - conversation_id: str = Path(..., description="Conversation ID"), - planning_context: Optional[Dict[str, Any]] = Body( - None, description="Additional planning context" - ), -): - """Generate planning milestones from conversation""" - try: - milestones = ( - await app.state.goal_conversation_service.generate_planning_milestones( - conversation_id=conversation_id, planning_context=planning_context - ) - ) - - return { - "conversation_id": conversation_id, - "milestones": milestones, - "count": len(milestones), - "status": "generated", - } - - except Exception as e: - logger.error(f"Error generating planning milestones: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/conversations/{conversation_id}/conduct-progress-review", - tags=["goals-management"], - summary="Conduct progress review", - description="Conduct AI-powered progress review for a goal conversation", -) -async def conduct_progress_review( - conversation_id: str = Path(..., description="Conversation ID"), - review_period_days: int = Query( - 30, ge=1, le=365, description="Review period in days" - ), -): - """Conduct progress review""" - try: - review_analysis = ( - await app.state.goal_conversation_service.conduct_progress_review( - conversation_id=conversation_id, review_period_days=review_period_days - ) - ) - - return review_analysis - - except Exception as e: - logger.error(f"Error conducting progress review: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/conversations/{conversation_id}/extract-action-items", - tags=["goals-management"], - summary="Extract action items", - description="Extract and create action items from conversation analysis", -) -async def extract_action_items( - conversation_id: str = Path(..., description="Conversation ID"), - auto_assign: bool = Query( - True, description="Whether to automatically assign action items" - ), -): - """Extract action items from conversation""" - try: - action_items = await app.state.goal_conversation_service.extract_action_items_from_conversation( - conversation_id=conversation_id, auto_assign=auto_assign - ) - - return { - "conversation_id": conversation_id, - "action_items": action_items, - "count": len(action_items), - "status": "extracted", - } - - except Exception as e: - logger.error(f"Error extracting action items: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/goals/{goal_id}/conversations", - tags=["goals-management"], - summary="Get goal conversations", - description="Get all conversations for a goal with optional filtering", -) -async def get_goal_conversations( - goal_id: str = Path(..., description="Goal ID"), - conversation_type: Optional[str] = Query( - None, description="Filter by conversation type" - ), - status: Optional[str] = Query(None, description="Filter by conversation status"), - limit: int = Query(10, ge=1, le=50, description="Maximum number of conversations"), -): - """Get conversations for a goal""" - try: - from .goal_conversation_service import ConversationStatus, ConversationType - - conv_type = ConversationType(conversation_type) if conversation_type else None - conv_status = ConversationStatus(status) if status else None - - conversations = ( - await app.state.goal_conversation_service.get_goal_conversations( - goal_id=goal_id, - conversation_type=conv_type, - status=conv_status, - limit=limit, - ) - ) - - return { - "goal_id": goal_id, - "conversations": conversations, - "count": len(conversations), - } - - except Exception as e: - logger.error(f"Error getting goal conversations: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/goals/{goal_id}/track-progress", - tags=["goals-management"], - summary="Record progress tracking update", - description="Record detailed progress update with tracking and risk assessment", -) -async def record_progress_tracking( - goal_id: str = Path(..., description="Goal ID"), - progress_data: ProgressUpdateRequest = Body( - ..., description="Progress tracking data" - ), - recorded_by: Optional[str] = Query( - None, description="ID of user/agent recording progress" - ), -): - """Record progress tracking update""" - try: - snapshot_id = await app.state.goal_tracking_service.record_progress_update( - goal_id=goal_id, - progress_percentage=progress_data.progress_percentage, - current_value=progress_data.current_value, - milestone_id=progress_data.milestone_id, - notes=progress_data.notes, - recorded_by=recorded_by, - confidence_score=progress_data.confidence_score, - trigger_alerts=progress_data.trigger_alerts, - ) - - return {"goal_id": goal_id, "snapshot_id": snapshot_id, "status": "recorded"} - - except Exception as e: - logger.error(f"Error recording progress tracking: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/goals/{goal_id}/deadline-risk", - tags=["goals-management"], - summary="Assess deadline risk", - description="Get comprehensive deadline risk assessment for a goal", -) -async def assess_deadline_risk(goal_id: str = Path(..., description="Goal ID")): - """Assess deadline risk for goal""" - try: - deadline_risk = await app.state.goal_tracking_service.assess_goal_deadline_risk( - goal_id - ) - - return { - "goal_id": deadline_risk.goal_id, - "risk_level": deadline_risk.risk_level.value, - "probability_of_delay": float(deadline_risk.probability_of_delay), - "estimated_completion_date": deadline_risk.estimated_completion_date.isoformat(), - "days_at_risk": deadline_risk.days_at_risk, - "critical_path_items": deadline_risk.critical_path_items, - "mitigation_strategies": deadline_risk.mitigation_strategies, - "updated_at": deadline_risk.updated_at.isoformat(), - } - - except Exception as e: - logger.error(f"Error assessing deadline risk: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/goals/{goal_id}/progress-report", - tags=["goals-management"], - summary="Generate progress report", - description="Generate comprehensive progress report for a goal", -) -async def generate_progress_report( - goal_id: str = Path(..., description="Goal ID"), - report_period_days: int = Query( - 30, ge=1, le=365, description="Report period in days" - ), -): - """Generate progress report for goal""" - try: - report = await app.state.goal_tracking_service.generate_progress_report( - goal_id=goal_id, report_period_days=report_period_days - ) - - return report - - except Exception as e: - logger.error(f"Error generating progress report: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/organizations/{organization_id}/goals-dashboard", - tags=["goals-management"], - summary="Get organization goals dashboard", - description="Get comprehensive dashboard for all organization goals", -) -async def get_organization_goals_dashboard( - organization_id: str = Path(..., description="Organization ID") -): - """Get organization goals dashboard""" - try: - dashboard = await app.state.goals_service.get_organization_goals_dashboard( - organization_id - ) - return dashboard - - except Exception as e: - logger.error(f"Error getting organization dashboard: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/organizations/{organization_id}/tracking-dashboard", - tags=["goals-management"], - summary="Get tracking dashboard", - description="Get comprehensive tracking dashboard with risk assessments", -) -async def get_tracking_dashboard( - organization_id: str = Path(..., description="Organization ID") -): - """Get organization tracking dashboard""" - try: - dashboard = ( - await app.state.goal_tracking_service.get_organization_tracking_dashboard( - organization_id - ) - ) - return dashboard - - except Exception as e: - logger.error(f"Error getting tracking dashboard: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# ================================ -# Knowledge Management API Endpoints -# ================================ - - -@app.post( - "/knowledge/organizations/{organization_id}/documents", - tags=["knowledge-management"], - summary="Upload Organizational Document", - response_model=DocumentMetadata, -) -async def upload_organization_document( - organization_id: str = Path(..., description="Organization ID"), - file: UploadFile = File(..., description="Document file to upload"), - title: Optional[str] = Form(None, description="Document title"), - tags: Optional[str] = Form(None, description="Comma-separated tags"), -): - """Upload a document to organizational knowledge base""" - try: - tags_list = [] - if tags: - tags_list = [tag.strip() for tag in tags.split(",")] - - document = await knowledge_manager.upload_document( - file_content=file.file, - filename=file.filename, - title=title, - organization_id=organization_id, - tags=tags_list, - ) - - return document - - except Exception as e: - logger.error(f"Error uploading organizational document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/knowledge/organizations/{organization_id}/url", - tags=["knowledge-management"], - summary="Add URL to Organizational Knowledge", - response_model=DocumentMetadata, -) -async def add_organization_url( - organization_id: str = Path(..., description="Organization ID"), - url: str = Body(..., embed=True), - title: Optional[str] = Body(None, embed=True), - tags: Optional[List[str]] = Body(None, embed=True), -): - """Add URL content to organizational knowledge base""" - try: - document = await knowledge_manager.upload_url( - url=url, title=title, organization_id=organization_id, tags=tags or [] - ) - - return document - - except Exception as e: - logger.error(f"Error adding organizational URL: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/organizations/{organization_id}/documents", - tags=["knowledge-management"], - summary="List Organizational Documents", - response_model=List[DocumentMetadata], -) -async def list_organization_documents( - organization_id: str = Path(..., description="Organization ID") -): - """Get list of organizational documents""" - try: - documents = await knowledge_manager.get_documents( - organization_id=organization_id - ) - return documents - - except Exception as e: - logger.error(f"Error listing organizational documents: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/organizations/{organization_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Get Organizational Document", - response_model=DocumentMetadata, -) -async def get_organization_document( - organization_id: str = Path(..., description="Organization ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Get organizational document metadata""" - try: - document = await knowledge_manager.get_document_metadata( - doc_id=doc_id, organization_id=organization_id - ) - - if not document: - raise HTTPException(status_code=404, detail="Document not found") - - return document - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting organizational document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/organizations/{organization_id}/documents/{doc_id}/content", - tags=["knowledge-management"], - summary="Get Organizational Document Content", -) -async def get_organization_document_content( - organization_id: str = Path(..., description="Organization ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Get full content of organizational document""" - try: - content = await knowledge_manager.get_document_content( - doc_id=doc_id, organization_id=organization_id - ) - - if content is None: - raise HTTPException(status_code=404, detail="Document not found") - - return {"content": content} - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting organizational document content: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.put( - "/knowledge/organizations/{organization_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Update Organizational Document", - response_model=DocumentMetadata, -) -async def update_organization_document( - organization_id: str = Path(..., description="Organization ID"), - doc_id: str = Path(..., description="Document ID"), - title: Optional[str] = Body(None, embed=True), - tags: Optional[List[str]] = Body(None, embed=True), -): - """Update organizational document metadata""" - try: - document = await knowledge_manager.update_document( - doc_id=doc_id, title=title, tags=tags, organization_id=organization_id - ) - - if not document: - raise HTTPException(status_code=404, detail="Document not found") - - return document - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error updating organizational document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.delete( - "/knowledge/organizations/{organization_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Delete Organizational Document", -) -async def delete_organization_document( - organization_id: str = Path(..., description="Organization ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Delete organizational document""" - try: - success = await knowledge_manager.delete_document( - doc_id=doc_id, organization_id=organization_id - ) - - if not success: - raise HTTPException(status_code=404, detail="Document not found") - - return {"message": "Document deleted successfully"} - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error deleting organizational document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# Team Knowledge Management Endpoints - - -@app.post( - "/knowledge/teams/{team_id}/documents", - tags=["knowledge-management"], - summary="Upload Team Document", - response_model=DocumentMetadata, -) -async def upload_team_document( - team_id: str = Path(..., description="Team ID"), - file: UploadFile = File(..., description="Document file to upload"), - title: Optional[str] = Form(None, description="Document title"), - tags: Optional[str] = Form(None, description="Comma-separated tags"), -): - """Upload a document to team knowledge base""" - try: - tags_list = [] - if tags: - tags_list = [tag.strip() for tag in tags.split(",")] - - document = await knowledge_manager.upload_document( - file_content=file.file, - filename=file.filename, - title=title, - team_id=team_id, - tags=tags_list, - ) - - return document - - except Exception as e: - logger.error(f"Error uploading team document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/knowledge/teams/{team_id}/url", - tags=["knowledge-management"], - summary="Add URL to Team Knowledge", - response_model=DocumentMetadata, -) -async def add_team_url( - team_id: str = Path(..., description="Team ID"), - url: str = Body(..., embed=True), - title: Optional[str] = Body(None, embed=True), - tags: Optional[List[str]] = Body(None, embed=True), -): - """Add URL content to team knowledge base""" - try: - document = await knowledge_manager.upload_url( - url=url, title=title, team_id=team_id, tags=tags or [] - ) - - return document - - except Exception as e: - logger.error(f"Error adding team URL: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/teams/{team_id}/documents", - tags=["knowledge-management"], - summary="List Team Documents", - response_model=List[DocumentMetadata], -) -async def list_team_documents(team_id: str = Path(..., description="Team ID")): - """Get list of team documents""" - try: - documents = await knowledge_manager.get_documents(team_id=team_id) - return documents - - except Exception as e: - logger.error(f"Error listing team documents: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/teams/{team_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Get Team Document", - response_model=DocumentMetadata, -) -async def get_team_document( - team_id: str = Path(..., description="Team ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Get team document metadata""" - try: - document = await knowledge_manager.get_document_metadata( - doc_id=doc_id, team_id=team_id - ) - - if not document: - raise HTTPException(status_code=404, detail="Document not found") - - return document - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting team document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/teams/{team_id}/documents/{doc_id}/content", - tags=["knowledge-management"], - summary="Get Team Document Content", -) -async def get_team_document_content( - team_id: str = Path(..., description="Team ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Get full content of team document""" - try: - content = await knowledge_manager.get_document_content( - doc_id=doc_id, team_id=team_id - ) - - if content is None: - raise HTTPException(status_code=404, detail="Document not found") - - return {"content": content} - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting team document content: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.put( - "/knowledge/teams/{team_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Update Team Document", - response_model=DocumentMetadata, -) -async def update_team_document( - team_id: str = Path(..., description="Team ID"), - doc_id: str = Path(..., description="Document ID"), - title: Optional[str] = Body(None, embed=True), - tags: Optional[List[str]] = Body(None, embed=True), -): - """Update team document metadata""" - try: - document = await knowledge_manager.update_document( - doc_id=doc_id, title=title, tags=tags, team_id=team_id - ) - - if not document: - raise HTTPException(status_code=404, detail="Document not found") - - return document - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error updating team document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.delete( - "/knowledge/teams/{team_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Delete Team Document", -) -async def delete_team_document( - team_id: str = Path(..., description="Team ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Delete team document""" - try: - success = await knowledge_manager.delete_document( - doc_id=doc_id, team_id=team_id - ) - - if not success: - raise HTTPException(status_code=404, detail="Document not found") - - return {"message": "Document deleted successfully"} - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error deleting team document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# Agent Knowledge Management Endpoints - - -@app.post( - "/knowledge/agents/{agent_id}/documents", - tags=["knowledge-management"], - summary="Upload Agent Document", - response_model=DocumentMetadata, -) -async def upload_agent_document( - agent_id: str = Path(..., description="Agent ID"), - file: UploadFile = File(..., description="Document file to upload"), - title: Optional[str] = Form(None, description="Document title"), - tags: Optional[str] = Form(None, description="Comma-separated tags"), -): - """Upload a document to agent knowledge base""" - try: - tags_list = [] - if tags: - tags_list = [tag.strip() for tag in tags.split(",")] - - document = await knowledge_manager.upload_document( - file_content=file.file, - filename=file.filename, - title=title, - agent_id=agent_id, - tags=tags_list, - ) - - return document - - except Exception as e: - logger.error(f"Error uploading agent document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/knowledge/agents/{agent_id}/url", - tags=["knowledge-management"], - summary="Add URL to Agent Knowledge", - response_model=DocumentMetadata, -) -async def add_agent_url( - agent_id: str = Path(..., description="Agent ID"), - url: str = Body(..., embed=True), - title: Optional[str] = Body(None, embed=True), - tags: Optional[List[str]] = Body(None, embed=True), -): - """Add URL content to agent knowledge base""" - try: - document = await knowledge_manager.upload_url( - url=url, title=title, agent_id=agent_id, tags=tags or [] - ) - - return document - - except Exception as e: - logger.error(f"Error adding agent URL: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/agents/{agent_id}/documents", - tags=["knowledge-management"], - summary="List Agent Documents", - response_model=List[DocumentMetadata], -) -async def list_agent_documents(agent_id: str = Path(..., description="Agent ID")): - """Get list of agent documents""" - try: - documents = await knowledge_manager.get_documents(agent_id=agent_id) - return documents - - except Exception as e: - logger.error(f"Error listing agent documents: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/agents/{agent_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Get Agent Document", - response_model=DocumentMetadata, -) -async def get_agent_document( - agent_id: str = Path(..., description="Agent ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Get agent document metadata""" - try: - document = await knowledge_manager.get_document_metadata( - doc_id=doc_id, agent_id=agent_id - ) - - if not document: - raise HTTPException(status_code=404, detail="Document not found") - - return document - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting agent document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/knowledge/agents/{agent_id}/documents/{doc_id}/content", - tags=["knowledge-management"], - summary="Get Agent Document Content", -) -async def get_agent_document_content( - agent_id: str = Path(..., description="Agent ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Get full content of agent document""" - try: - content = await knowledge_manager.get_document_content( - doc_id=doc_id, agent_id=agent_id - ) - - if content is None: - raise HTTPException(status_code=404, detail="Document not found") - - return {"content": content} - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error getting agent document content: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.put( - "/knowledge/agents/{agent_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Update Agent Document", - response_model=DocumentMetadata, -) -async def update_agent_document( - agent_id: str = Path(..., description="Agent ID"), - doc_id: str = Path(..., description="Document ID"), - title: Optional[str] = Body(None, embed=True), - tags: Optional[List[str]] = Body(None, embed=True), -): - """Update agent document metadata""" - try: - document = await knowledge_manager.update_document( - doc_id=doc_id, title=title, tags=tags, agent_id=agent_id - ) - - if not document: - raise HTTPException(status_code=404, detail="Document not found") - - return document - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error updating agent document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.delete( - "/knowledge/agents/{agent_id}/documents/{doc_id}", - tags=["knowledge-management"], - summary="Delete Agent Document", -) -async def delete_agent_document( - agent_id: str = Path(..., description="Agent ID"), - doc_id: str = Path(..., description="Document ID"), -): - """Delete agent document""" - try: - success = await knowledge_manager.delete_document( - doc_id=doc_id, agent_id=agent_id - ) - - if not success: - raise HTTPException(status_code=404, detail="Document not found") - - return {"message": "Document deleted successfully"} - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error deleting agent document: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# Knowledge Search Endpoints - - -@app.get( - "/knowledge/search", - tags=["knowledge-management"], - summary="Search Knowledge Base", - response_model=List[DocumentMetadata], -) -async def search_knowledge( - query: str = Query(..., description="Search query"), - organization_id: Optional[str] = Query(None, description="Filter by organization"), - team_id: Optional[str] = Query(None, description="Filter by team"), - agent_id: Optional[str] = Query(None, description="Filter by agent"), - limit: int = Query(10, ge=1, le=100, description="Maximum number of results"), -): - """Search across knowledge base""" - try: - documents = await knowledge_manager.search_documents( - query=query, - organization_id=organization_id, - team_id=team_id, - agent_id=agent_id, - limit=limit, - ) - - return documents - - except Exception as e: - logger.error(f"Error searching knowledge: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# ================================ -# Container Management API Endpoints -# ================================ - - -@app.post( - "/agents/{agent_id}/container/create", - tags=["container-management"], - summary="Create Agent Container", - response_model=ContainerStatus, -) -async def create_agent_container( - agent_id: str = Path(..., description="Agent ID"), - config: Optional[ContainerConfig] = Body( - None, description="Container configuration" - ), -): - """Create a new container for an AI agent""" - try: - status = await container_manager.create_agent_container(agent_id, config) - return status - - except Exception as e: - logger.error(f"Error creating container for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/agents/{agent_id}/container/start", - tags=["container-management"], - summary="Start Agent Container", - response_model=ContainerStatus, -) -async def start_agent_container(agent_id: str = Path(..., description="Agent ID")): - """Start an agent container""" - try: - status = await container_manager.start_container(agent_id) - return status - - except RuntimeError as e: - raise HTTPException(status_code=404, detail=str(e)) - except Exception as e: - logger.error(f"Error starting container for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/agents/{agent_id}/container/stop", - tags=["container-management"], - summary="Stop Agent Container", - response_model=ContainerStatus, -) -async def stop_agent_container( - agent_id: str = Path(..., description="Agent ID"), - timeout: int = Body(30, description="Stop timeout in seconds"), -): - """Stop an agent container""" - try: - status = await container_manager.stop_container(agent_id, timeout) - return status - - except RuntimeError as e: - raise HTTPException(status_code=404, detail=str(e)) - except Exception as e: - logger.error(f"Error stopping container for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/agents/{agent_id}/container/restart", - tags=["container-management"], - summary="Restart Agent Container", - response_model=ContainerStatus, -) -async def restart_agent_container( - agent_id: str = Path(..., description="Agent ID"), - timeout: int = Body(30, description="Restart timeout in seconds"), -): - """Restart an agent container""" - try: - status = await container_manager.restart_container(agent_id, timeout) - return status - - except RuntimeError as e: - raise HTTPException(status_code=404, detail=str(e)) - except Exception as e: - logger.error(f"Error restarting container for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.delete( - "/agents/{agent_id}/container", - tags=["container-management"], - summary="Remove Agent Container", -) -async def remove_agent_container( - agent_id: str = Path(..., description="Agent ID"), - force: bool = Query(False, description="Force removal of running container"), -): - """Remove an agent container""" - try: - success = await container_manager.remove_container(agent_id, force) - - if success: - return {"message": f"Container for agent {agent_id} removed successfully"} - else: - raise HTTPException(status_code=500, detail="Failed to remove container") - - except Exception as e: - logger.error(f"Error removing container for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/agents/{agent_id}/container/status", - tags=["container-management"], - summary="Get Agent Container Status", - response_model=Optional[ContainerStatus], -) -async def get_agent_container_status(agent_id: str = Path(..., description="Agent ID")): - """Get container status for an agent""" - try: - status = await container_manager.get_container_status(agent_id) - return status - - except Exception as e: - logger.error(f"Error getting container status for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/containers/agents", - tags=["container-management"], - summary="List Agent Containers", - response_model=List[ContainerStatus], -) -async def list_agent_containers(): - """List all agent containers""" - try: - containers = await container_manager.list_agent_containers() - return containers - - except Exception as e: - logger.error(f"Error listing agent containers: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get( - "/agents/{agent_id}/container/logs", - tags=["container-management"], - summary="Get Agent Container Logs", -) -async def get_agent_container_logs( - agent_id: str = Path(..., description="Agent ID"), - tail: int = Query(100, ge=1, le=10000, description="Number of log lines to return"), - since: Optional[str] = Query( - None, description="Show logs since timestamp (ISO format)" - ), -): - """Get container logs for an agent""" - try: - since_dt = None - if since: - try: - since_dt = datetime.fromisoformat(since.replace("Z", "+00:00")) - except ValueError: - raise HTTPException(status_code=400, detail="Invalid timestamp format") - - logs = await container_manager.get_container_logs( - agent_id=agent_id, tail=tail, since=since_dt - ) - - return {"logs": logs} - - except RuntimeError as e: - raise HTTPException(status_code=404, detail=str(e)) - except Exception as e: - logger.error(f"Error getting container logs for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post( - "/agents/{agent_id}/container/execute", - tags=["container-management"], - summary="Execute Command in Container", -) -async def execute_container_command( - agent_id: str = Path(..., description="Agent ID"), - command: str = Body(..., description="Command to execute"), - working_dir: Optional[str] = Body(None, description="Working directory"), -): - """Execute a command in the agent container""" - try: - result = await container_manager.execute_command( - agent_id=agent_id, command=command, working_dir=working_dir - ) - - return result - - except RuntimeError as e: - raise HTTPException(status_code=404, detail=str(e)) - except Exception as e: - logger.error(f"Error executing command in container for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# WebSocket endpoint for real-time log streaming -@app.websocket("/agents/{agent_id}/container/logs/stream") -async def stream_agent_container_logs(websocket: WebSocket, agent_id: str): - """Stream container logs in real-time via WebSocket""" - await websocket.accept() - - try: - # Check if container exists - status = await container_manager.get_container_status(agent_id) - if not status: - await websocket.send_json({"error": "Container not found"}) - await websocket.close() - return - - await websocket.send_json({"status": "connected", "agent_id": agent_id}) - - # Stream logs - async for log_entry in container_manager.stream_container_logs(agent_id): - await websocket.send_json( - { - "timestamp": log_entry.timestamp.isoformat(), - "stream": log_entry.stream, - "message": log_entry.message, - } - ) - - except Exception as e: - logger.error(f"Error in log stream for agent {agent_id}: {e}") - try: - await websocket.send_json({"error": str(e)}) - except Exception: - logger.debug("Failed to send error frame on closing websocket") - finally: - try: - await websocket.close() - except Exception: - logger.debug("Failed to close websocket cleanly") - - -# ============================================================================ -# RAG (Retrieval-Augmented Generation) Endpoints -# ============================================================================ - - -@app.post("/rag/search") -async def search_knowledge_context( - query: str = Body(..., embed=True), - organization_id: Optional[str] = Body(None, embed=True), - team_id: Optional[str] = Body(None, embed=True), - agent_id: Optional[str] = Body(None, embed=True), - max_results: int = Body(5, embed=True), - similarity_threshold: float = Body(0.7, embed=True), -): - """Search for relevant knowledge context using RAG""" - try: - context = await rag_system.search_relevant_context( - query=query, - organization_id=organization_id, - team_id=team_id, - agent_id=agent_id, - max_results=max_results, - similarity_threshold=similarity_threshold, - ) - - return { - "query": context.query, - "relevant_chunks": [ - { - "document_id": chunk.document_id, - "document_title": chunk.metadata.get("document_title", "Unknown"), - "content": chunk.content, - "chunk_index": chunk.chunk_index, - "metadata": chunk.metadata, - } - for chunk in context.relevant_chunks - ], - "similarity_scores": context.similarity_scores, - "total_documents": context.total_documents, - "context_length": context.context_length, - } - - except Exception as e: - logger.error(f"Error searching knowledge context: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post("/rag/enhance-prompt") -async def enhance_prompt_with_context( - message: str = Body(..., embed=True), - organization_id: Optional[str] = Body(None, embed=True), - team_id: Optional[str] = Body(None, embed=True), - agent_id: Optional[str] = Body(None, embed=True), - max_context_length: int = Body(4000, embed=True), -): - """Enhance a prompt with relevant context using RAG""" - try: - enhanced_prompt = await rag_system.get_contextual_prompt( - user_message=message, - organization_id=organization_id, - team_id=team_id, - agent_id=agent_id, - max_context_length=max_context_length, - ) - - return { - "original_message": message, - "enhanced_prompt": enhanced_prompt, - "context_added": len(enhanced_prompt) > len(message), - } - - except Exception as e: - logger.error(f"Error enhancing prompt with context: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post("/rag/reindex") -async def reindex_knowledge_base( - organization_id: Optional[str] = Body(None, embed=True), - team_id: Optional[str] = Body(None, embed=True), - agent_id: Optional[str] = Body(None, embed=True), -): - """Reindex all documents in a scope for RAG""" - try: - results = await rag_system.index_all_documents( - organization_id=organization_id, team_id=team_id, agent_id=agent_id - ) - - return { - "scope": { - "organization_id": organization_id, - "team_id": team_id, - "agent_id": agent_id, - }, - "results": results, - "message": f"Indexed {results['indexed']} documents, {results['failed']} failed, {results['skipped']} skipped", - } - - except Exception as e: - logger.error(f"Error reindexing knowledge base: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.get("/rag/stats") -async def get_rag_index_stats(): - """Get statistics about the RAG index""" - try: - stats = await rag_system.get_index_stats() - return stats - - except Exception as e: - logger.error(f"Error getting RAG stats: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -@app.post("/rag/documents/{doc_id}/reindex") -async def reindex_document( - doc_id: str, - organization_id: Optional[str] = Body(None, embed=True), - team_id: Optional[str] = Body(None, embed=True), - agent_id: Optional[str] = Body(None, embed=True), -): - """Reindex a specific document for RAG""" - try: - # Get document metadata - document = await knowledge_manager.get_document_metadata( - doc_id=doc_id, - organization_id=organization_id, - team_id=team_id, - agent_id=agent_id, - ) - - if not document: - raise HTTPException(status_code=404, detail="Document not found") - - # Reindex the document - success = await rag_system.index_document(document) - - if success: - return { - "document_id": doc_id, - "status": "reindexed", - "message": f"Document '{document.title}' has been reindexed successfully", - } - else: - raise HTTPException(status_code=500, detail="Failed to reindex document") - - except HTTPException: - raise - except Exception as e: - logger.error(f"Error reindexing document {doc_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# ============================================================================ -# Real-time WebSocket Endpoints -# ============================================================================ - - -@app.websocket("/ws/updates") -async def websocket_real_time_updates( - websocket: WebSocket, - organization_id: Optional[str] = None, - team_id: Optional[str] = None, - agent_id: Optional[str] = None, - user_id: Optional[str] = None, - subscriptions: Optional[str] = None, -): - """Main WebSocket endpoint for real-time updates""" - import uuid - - connection_id = str(uuid.uuid4()) - - # Parse subscriptions - subscription_list = [] - if subscriptions: - subscription_list = subscriptions.split(",") - - try: - connection = await websocket_manager.connect( - websocket=websocket, - connection_id=connection_id, - organization_id=organization_id, - team_id=team_id, - agent_id=agent_id, - user_id=user_id, - subscriptions=subscription_list, - ) - - # Keep connection alive and handle pings - while True: - try: - # Wait for ping messages or disconnection - message = await websocket.receive_text() - - # Handle ping/pong - if message == "ping": - await websocket.send_text("pong") - connection.last_ping = datetime.now() - else: - # Parse other messages (subscription updates, etc.) - try: - data = json.loads(message) - if data.get("type") == "subscribe": - # Update subscriptions - new_subs = data.get("subscriptions", []) - connection.scope.subscriptions.clear() - for sub in new_subs: - try: - connection.scope.subscriptions.add(UpdateType(sub)) - except ValueError: - pass - - await connection.send_update( - WebSocketUpdate( - type=UpdateType.SYSTEM_NOTIFICATION, - data={ - "message": "Subscriptions updated", - "subscriptions": list( - connection.scope.subscriptions - ), - }, - ) - ) - except json.JSONDecodeError: - pass - - except WebSocketDisconnect: - break - - except Exception as e: - logger.error(f"WebSocket error for connection {connection_id}: {e}") - finally: - await websocket_manager.disconnect(connection_id) - - -@app.websocket("/ws/agent/{agent_id}/updates") -async def websocket_agent_updates(websocket: WebSocket, agent_id: str): - """WebSocket endpoint for specific agent updates""" - import uuid - - connection_id = f"agent-{agent_id}-{uuid.uuid4()}" - - try: - connection = await websocket_manager.connect( - websocket=websocket, - connection_id=connection_id, - agent_id=agent_id, - subscriptions=[ - UpdateType.AGENT_STATUS.value, - UpdateType.TASK_STATUS.value, - UpdateType.TASK_PROGRESS.value, - UpdateType.CONTAINER_STATUS.value, - UpdateType.CHAT_MESSAGE.value, - UpdateType.CHAT_TYPING.value, - ], - ) - - # Keep connection alive - while True: - try: - message = await websocket.receive_text() - if message == "ping": - await websocket.send_text("pong") - connection.last_ping = datetime.now() - except WebSocketDisconnect: - break - - except Exception as e: - logger.error(f"Agent WebSocket error for {agent_id}: {e}") - finally: - await websocket_manager.disconnect(connection_id) - - -@app.websocket("/ws/agents/{agent_id}/conversations/{conversation_id}") -async def websocket_agent_conversation( - websocket: WebSocket, agent_id: str, conversation_id: str -): - """WebSocket endpoint for real-time agent conversation""" - import uuid - - connection_id = f"conversation-{conversation_id}-{uuid.uuid4()}" - - await websocket.accept() - - try: - # Store connection for broadcasting - active_conversations = getattr(app.state, "active_conversations", {}) - if conversation_id not in active_conversations: - active_conversations[conversation_id] = [] - active_conversations[conversation_id].append(websocket) - app.state.active_conversations = active_conversations - - while True: - try: - # Receive message from client - data = await websocket.receive_json() - - if data.get("type") == "ping": - await websocket.send_json({"type": "pong"}) - elif data.get("type") == "message": - # Handle new message - message_content = data.get("content", "") - if message_content: - # Store message in database - async with get_db_connection() as conn: - message_id = await conn.fetchval( - """ - INSERT INTO agent_conversations (session_id, agent_id, message_type, content) - VALUES ($1, $2, 'user', $3) - RETURNING id - """, - conversation_id, - agent_id, - message_content, - ) - - # Update session activity - await conn.execute( - """ - UPDATE chat_sessions - SET last_activity = CURRENT_TIMESTAMP, message_count = message_count + 1 - WHERE id = $1 - """, - conversation_id, - ) - - # Broadcast to all connected clients for this conversation - message_data = { - "type": "new_message", - "message": { - "id": str(message_id), - "conversation_id": conversation_id, - "role": "user", - "content": message_content, - "timestamp": datetime.now().isoformat(), - "status": "sent", - }, - } - - for conn in active_conversations.get(conversation_id, []): - try: - await conn.send_json(message_data) - except Exception: - # Connection might be closed; skip this subscriber - logger.debug("Skipped broadcast to a closed websocket") - - # TODO: Here we would trigger agent response generation - # For now, send a simple acknowledgment after a delay - await asyncio.sleep(1) - - agent_response = { - "type": "new_message", - "message": { - "id": str(uuid.uuid4()), - "conversation_id": conversation_id, - "role": "agent", - "content": f"I received your message: {message_content}", - "timestamp": datetime.now().isoformat(), - "status": "received", - }, - } - - for conn in active_conversations.get(conversation_id, []): - try: - await conn.send_json(agent_response) - except Exception: - # Connection might be closed; skip this subscriber - logger.debug("Skipped broadcast to a closed websocket") - - # Store agent response in database - async with get_db_connection() as conn: - await conn.execute( - """ - INSERT INTO agent_conversations (session_id, agent_id, message_type, content) - VALUES ($1, $2, 'agent', $3) - """, - conversation_id, - agent_id, - agent_response["message"]["content"], - ) - - except WebSocketDisconnect: - break - except Exception as e: - logger.error(f"Error in conversation WebSocket: {e}") - - except Exception as e: - logger.error(f"Conversation WebSocket error for {conversation_id}: {e}") - finally: - # Clean up connection - if ( - hasattr(app.state, "active_conversations") - and conversation_id in app.state.active_conversations - ): - if websocket in app.state.active_conversations[conversation_id]: - app.state.active_conversations[conversation_id].remove(websocket) - - -@app.websocket("/ws/organization/{organization_id}/updates") -async def websocket_organization_updates(websocket: WebSocket, organization_id: str): - """WebSocket endpoint for organization-wide updates""" - import uuid - - connection_id = f"org-{organization_id}-{uuid.uuid4()}" - - try: - connection = await websocket_manager.connect( - websocket=websocket, - connection_id=connection_id, - organization_id=organization_id, - subscriptions=[ - UpdateType.AGENT_CREATED.value, - UpdateType.AGENT_UPDATED.value, - UpdateType.AGENT_DELETED.value, - UpdateType.KNOWLEDGE_UPDATED.value, - UpdateType.KNOWLEDGE_INDEXED.value, - UpdateType.SYSTEM_NOTIFICATION.value, - ], - ) - - # Keep connection alive - while True: - try: - message = await websocket.receive_text() - if message == "ping": - await websocket.send_text("pong") - connection.last_ping = datetime.now() - except WebSocketDisconnect: - break - - except Exception as e: - logger.error(f"Organization WebSocket error for {organization_id}: {e}") - finally: - await websocket_manager.disconnect(connection_id) - - -# WebSocket Statistics Endpoint -@app.get("/ws/stats") -async def get_websocket_stats(): - """Get WebSocket connection statistics""" - try: - stats = websocket_manager.get_stats() - return stats - except Exception as e: - logger.error(f"Error getting WebSocket stats: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# Manual notification endpoints for testing -@app.post("/ws/test/agent/{agent_id}/status") -async def test_agent_status_notification( - agent_id: str, - status: str = Body(..., embed=True), - message: Optional[str] = Body(None, embed=True), -): - """Test endpoint to send agent status notifications""" - try: - await notify_agent_status_change( - agent_id=agent_id, - status=status, - additional_data={"message": message} if message else None, - ) - return {"status": "notification_sent", "agent_id": agent_id} - except Exception as e: - logger.error(f"Error sending test notification: {e}") - raise HTTPException(status_code=500, detail=str(e)) - - -# ============================================================================ -# Missing API Endpoints (Goals, Teams, Organizations) -# ============================================================================ - - -@app.get("/teams") -async def get_teams(): - """Get list of teams""" - # Mock data for now - return [ - { - "id": "1", - "name": "Development Team", - "description": "Frontend and backend developers", - "member_count": 5, - "organization_id": "1", - }, - { - "id": "2", - "name": "Executive Team", - "description": "Leadership and strategy", - "member_count": 3, - "organization_id": "1", - }, - ] - - -@app.get("/organizations/{organization_id}/goals") -async def get_organization_goals(organization_id: str): - """Get goals for an organization""" - # Mock data for now - return [ - { - "id": "1", - "title": "Increase Development Velocity", - "description": "Improve team productivity and code quality", - "status": "active", - "progress": 75, - "organization_id": organization_id, - "created_at": "2024-01-15T10:00:00Z", - "due_date": "2024-12-31T23:59:59Z", - }, - { - "id": "2", - "title": "Enhance AI Capabilities", - "description": "Expand AI agent capabilities and intelligence", - "status": "active", - "progress": 50, - "organization_id": organization_id, - "created_at": "2024-02-01T10:00:00Z", - "due_date": "2024-11-30T23:59:59Z", - }, - ] - - -@app.get("/goals/{goal_id}") -async def get_goal_details(goal_id: str): - """Get detailed information about a specific goal""" - # Mock data for now - return { - "id": goal_id, - "title": "Increase Development Velocity", - "description": "Improve team productivity and code quality through better tooling, processes, and automation", - "status": "active", - "progress": 75, - "organization_id": "1", - "team_id": "1", - "created_at": "2024-01-15T10:00:00Z", - "updated_at": "2024-08-06T16:30:00Z", - "due_date": "2024-12-31T23:59:59Z", - "milestones": [ - { - "id": "1", - "title": "Implement CI/CD Pipeline", - "description": "Set up automated testing and deployment", - "status": "completed", - "progress": 100, - "due_date": "2024-03-15T23:59:59Z", - }, - { - "id": "2", - "title": "Enhance Code Review Process", - "description": "Streamline code review workflow with automated tools", - "status": "in_progress", - "progress": 80, - "due_date": "2024-09-30T23:59:59Z", - }, - { - "id": "3", - "title": "Deploy AI-Powered Testing", - "description": "Implement intelligent test generation and execution", - "status": "planned", - "progress": 25, - "due_date": "2024-12-15T23:59:59Z", - }, - ], - "metrics": { - "deployment_frequency": "Daily", - "lead_time": "2.3 days", - "mttr": "45 minutes", - "change_failure_rate": "5%", - }, - "assigned_agents": [ - {"id": "1", "name": "DevOps Agent", "role": "CI/CD Specialist"}, - {"id": "2", "name": "QA Agent", "role": "Test Automation Engineer"}, - ], - } - - -@app.get("/agents/{agent_id}/tasks") -async def get_agent_tasks_list(agent_id: str): - """Get tasks for a specific agent - GET method""" - try: - # Get tasks from task queue - tasks = await app.state.task_queue.get_agent_tasks(agent_id) - return {"agent_id": agent_id, "tasks": tasks} - except Exception as e: - logger.error(f"Error getting tasks for agent {agent_id}: {e}") - raise HTTPException(status_code=500, detail=str(e)) +import asyncio +import json +import logging +import os +from collections import defaultdict +from contextlib import asynccontextmanager +from datetime import date, datetime +from decimal import Decimal +from typing import Any, Dict, List, Optional + +import jwt +from fastapi import ( + Body, + Depends, + FastAPI, + File, + Form, + HTTPException, + Path, + Query, + UploadFile, + WebSocket, + WebSocketDisconnect, + status, +) +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import HTMLResponse, Response +from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer +from pydantic import BaseModel, Field + +from hierarchy_endpoints import router as hierarchy_router + +from .agent_manager import AgentManager +from .container_manager import ContainerConfig, ContainerStatus, container_manager +from .context_service import ContextService +from .database import get_db_connection +from .knowledge_manager import DocumentMetadata, knowledge_manager +from .rag_integration import RAGContext, rag_system +from .sandbox_manager import AgentSandboxManager +from .task_execution_engine import TaskExecutionEngine +from .task_queue import TaskQueue +from .websocket_manager import ( + UpdateType, + WebSocketUpdate, + notify_agent_status_change, + notify_container_status_change, + notify_knowledge_update, + notify_task_progress, + websocket_manager, +) + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Auth helpers (Track 3) +# --------------------------------------------------------------------------- +_security = HTTPBearer(auto_error=False) +_jwt_secret = os.environ.get("FUZEFRONT_JWT_SECRET", "") + + +def require_auth(credentials: HTTPAuthorizationCredentials = Depends(_security)): + """Verify FuzeFront JWT on mutating endpoints. Disabled when secret not set (dev).""" + if not _jwt_secret: + return None # Auth disabled when secret not configured (dev mode) + if not credentials: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing token" + ) + try: + payload = jwt.decode(credentials.credentials, _jwt_secret, algorithms=["HS256"]) + return payload + except jwt.ExpiredSignatureError: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, detail="Token expired" + ) + except jwt.InvalidTokenError: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token" + ) + + +# --------------------------------------------------------------------------- +# Agent relay state (Track 4) +# --------------------------------------------------------------------------- +# agent_id -> list of subscriber WebSockets watching that agent's session +agent_relay_subscribers: Dict[str, List[WebSocket]] = defaultdict(list) + + +# Pydantic models for API documentation +class AgentCreateRequest(BaseModel): + name: str = Field(..., description="Agent name") + role: str = Field(..., description="Agent role (e.g., 'Senior React Developer')") + type: str = Field(..., description="Agent type (e.g., 'developer', 'executive')") + config: Dict[str, Any] = Field( + default_factory=dict, description="Agent configuration" + ) + repository_settings: Dict[str, Any] = Field( + default_factory=dict, description="Repository settings" + ) + sandbox_settings: Dict[str, Any] = Field( + default_factory=dict, description="Sandbox settings" + ) + + +class TaskCreateRequest(BaseModel): + title: str = Field(..., description="Task title") + description: str = Field(..., description="Task description") + priority: str = Field( + default="medium", description="Task priority (low, medium, high)" + ) + metadata: Dict[str, Any] = Field( + default_factory=dict, description="Additional task metadata" + ) + + +class HumanResponseRequest(BaseModel): + response: str = Field(..., description="Human response to agent question") + + +class FileOperationApprovalRequest(BaseModel): + approved: bool = Field(..., description="Whether to approve the file operations") + reason: Optional[str] = Field( + None, description="Optional reason for approval/rejection" + ) + + +class ClaudeSessionInputRequest(BaseModel): + input: str = Field(..., description="Input to send to Claude SDK session") + + +class CoordinationRequest(BaseModel): + coordination_mode: str = Field( + default="collaborative", + description="Coordination mode (sequential, parallel, hierarchical, collaborative)", + ) + required_agents: Optional[List[str]] = Field( + None, description="Specific agents to include" + ) + required_skills: Optional[List[str]] = Field( + None, description="Required skills for the task" + ) + + +class AgentCommunicationRequest(BaseModel): + message_type: str = Field( + default="notification", + description="Message type (request, response, notification, question)", + ) + content: str = Field(..., description="Message content") + metadata: Dict[str, Any] = Field( + default_factory=dict, description="Additional metadata" + ) + + +class MCPToolRequest(BaseModel): + tool_name: str = Field(..., description="Name of the MCP tool to call") + arguments: Dict[str, Any] = Field( + default_factory=dict, description="Tool arguments" + ) + + +class AgentMCPSetupRequest(BaseModel): + task_id: str = Field(..., description="Task ID for MCP setup") + session_id: Optional[str] = Field(None, description="Optional session ID") + + +class ConversationCreateRequest(BaseModel): + title: str = "New Conversation" + initial_message: Optional[str] = None + context: Optional[Dict[str, Any]] = None + + +class ConversationMessage(BaseModel): + role: str # 'user' or 'agent' + content: str + metadata: Optional[Dict[str, Any]] = None + + +class ChatMessageRequest(BaseModel): + content: str + metadata: Optional[Dict[str, Any]] = None + + +# Model Configuration Models +class ProviderCredentialsRequest(BaseModel): + provider: str = Field( + ..., description="Model provider (anthropic, openai, google, etc.)" + ) + api_key: str = Field(..., description="API key for the provider") + endpoint_url: Optional[str] = Field(None, description="Custom endpoint URL") + additional_config: Dict[str, Any] = Field( + default_factory=dict, description="Additional provider configuration" + ) + + +class AgentModelConfigRequest(BaseModel): + primary_model: str = Field(..., description="Primary model ID") + fallback_models: List[str] = Field( + default_factory=list, description="Fallback model IDs" + ) + temperature: float = Field( + default=0.7, ge=0.0, le=2.0, description="Model temperature" + ) + max_tokens: int = Field( + default=4096, ge=1, le=200000, description="Maximum output tokens" + ) + top_p: float = Field(default=1.0, ge=0.0, le=1.0, description="Top-p sampling") + frequency_penalty: float = Field( + default=0.0, ge=-2.0, le=2.0, description="Frequency penalty" + ) + presence_penalty: float = Field( + default=0.0, ge=-2.0, le=2.0, description="Presence penalty" + ) + custom_instructions: str = Field( + default="", description="Custom instructions for the agent" + ) + use_function_calling: bool = Field( + default=True, description="Enable function calling" + ) + streaming_enabled: bool = Field( + default=True, description="Enable response streaming" + ) + cost_limit_per_task: Optional[float] = Field( + None, ge=0.0, description="Cost limit per task in USD" + ) + + +class TaskCostEstimateRequest(BaseModel): + task_description: str = Field(..., description="Description of the task") + estimated_complexity: str = Field( + default="medium", + description="Estimated complexity (low, medium, high, very_high)", + ) + + +# Response models +class AgentResponse(BaseModel): + agent_id: str + status: str + agent: Dict[str, Any] + + +class TaskResponse(BaseModel): + task_id: str + status: str + + +class CoordinationResponse(BaseModel): + task_id: str + coordination_session_id: Optional[str] = None + status: str + coordination_mode: Optional[str] = None + message: Optional[str] = None + + +# Goals Management API Models +class GoalCreateRequest(BaseModel): + title: str = Field(..., description="Goal title") + description: str = Field(..., description="Goal description") + goal_type: str = Field( + default="business", + description="Goal type (business, technical, growth, operational)", + ) + target_value: Optional[Decimal] = Field( + None, description="Target value (e.g., 100000 for $100K)" + ) + target_unit: Optional[str] = Field( + None, description="Target unit (e.g., 'USD', 'users', '%')" + ) + target_deadline: Optional[date] = Field(None, description="Target completion date") + priority_level: int = Field( + default=5, ge=1, le=10, description="Priority level (1-10)" + ) + success_criteria: Optional[Dict[str, Any]] = Field( + default=None, description="Success criteria" + ) + assigned_teams: Optional[List[str]] = Field( + default=None, description="Assigned team IDs" + ) + goal_owner_agent_id: Optional[str] = Field(None, description="Goal owner agent ID") + stakeholder_agents: Optional[List[str]] = Field( + default=None, description="Stakeholder agent IDs" + ) + tags: Optional[List[str]] = Field(default=None, description="Goal tags") + metadata: Optional[Dict[str, Any]] = Field( + default=None, description="Additional metadata" + ) + + +class GoalUpdateRequest(BaseModel): + progress_percentage: Optional[Decimal] = Field( + None, ge=0, le=100, description="Progress percentage" + ) + current_value: Optional[Decimal] = Field(None, description="Current value") + completion_confidence: Optional[Decimal] = Field( + None, ge=0, le=1, description="Completion confidence" + ) + notes: Optional[str] = Field(None, description="Progress notes") + + +class MilestoneCreateRequest(BaseModel): + title: str = Field(..., description="Milestone title") + description: str = Field(..., description="Milestone description") + target_date: date = Field(..., description="Target completion date") + milestone_type: str = Field(default="deliverable", description="Milestone type") + success_criteria: Optional[Dict[str, Any]] = Field( + default=None, description="Success criteria" + ) + deliverables: Optional[List[Dict[str, Any]]] = Field( + default=None, description="Expected deliverables" + ) + dependencies: Optional[List[Dict[str, Any]]] = Field( + default=None, description="Dependencies" + ) + assigned_teams: Optional[List[str]] = Field( + default=None, description="Assigned teams" + ) + responsible_agent_id: Optional[str] = Field(None, description="Responsible agent") + priority_level: int = Field(default=5, ge=1, le=10, description="Priority level") + weight_in_goal: Optional[Decimal] = Field( + None, ge=0, le=100, description="Weight in goal (%)" + ) + + +class TaskFromMilestoneRequest(BaseModel): + title: str = Field(..., description="Task title") + description: str = Field(..., description="Task description") + task_type: str = Field(default="development", description="Task type") + complexity_level: str = Field(default="medium", description="Complexity level") + estimated_hours: Optional[Decimal] = Field(None, description="Estimated hours") + due_date: Optional[date] = Field(None, description="Due date") + assigned_team_id: Optional[str] = Field(None, description="Assigned team ID") + assigned_agent_id: Optional[str] = Field(None, description="Assigned agent ID") + priority: int = Field(default=5, ge=1, le=10, description="Priority") + requirements: Optional[Dict[str, Any]] = Field( + default=None, description="Requirements" + ) + acceptance_criteria: Optional[List[Dict[str, Any]]] = Field( + default=None, description="Acceptance criteria" + ) + dependencies: Optional[List[Dict[str, Any]]] = Field( + default=None, description="Dependencies" + ) + + +class GoalConversationCreateRequest(BaseModel): + conversation_type: str = Field(default="planning", description="Conversation type") + conversation_title: str = Field(..., description="Conversation title") + initial_context: Optional[Dict[str, Any]] = Field( + default=None, description="Initial context" + ) + participants: Optional[List[Dict[str, Any]]] = Field( + default=None, description="Participants" + ) + + +class ConversationMessageRequest(BaseModel): + message_type: str = Field(default="human", description="Message type") + sender_name: str = Field(..., description="Sender name") + content: str = Field(..., description="Message content") + metadata: Optional[Dict[str, Any]] = Field( + default=None, description="Message metadata" + ) + references: Optional[List[str]] = Field( + default=None, description="Referenced message IDs" + ) + + +class ProgressUpdateRequest(BaseModel): + progress_percentage: Optional[Decimal] = Field( + None, ge=0, le=100, description="Progress percentage" + ) + current_value: Optional[Decimal] = Field(None, description="Current value") + milestone_id: Optional[str] = Field(None, description="Associated milestone ID") + notes: Optional[str] = Field(None, description="Progress notes") + confidence_score: Optional[Decimal] = Field( + None, ge=0, le=1, description="Confidence score" + ) + trigger_alerts: bool = Field(default=True, description="Whether to trigger alerts") + + +@asynccontextmanager +async def lifespan(app: FastAPI): + # Startup + database_url = os.getenv( + "DATABASE_URL", "postgresql://postgres:password@postgres:5432/ai_context" + ) + + app.state.agent_manager = AgentManager(database_url) + app.state.task_queue = TaskQueue() + app.state.context_service = ContextService() + + # Initialize sandbox manager + app.state.sandbox_manager = AgentSandboxManager(database_url) + await app.state.sandbox_manager.start() + + # Start WebSocket manager background cleanup task + await websocket_manager.start() + + # Initialize task execution engine + app.state.task_execution_engine = TaskExecutionEngine(app.state.sandbox_manager) + await app.state.task_execution_engine.start() + + # Initialize multi-agent coordinator + from .multi_agent_coordinator import integrate_multi_agent_coordination + + app.state.multi_agent_coordinator = integrate_multi_agent_coordination( + app.state.task_execution_engine + ) + await app.state.multi_agent_coordinator.start() + + # Initialize knowledge management system + try: + from .context_enhancement_service import ContextEnhancementService + from .knowledge_notification_service import KnowledgeNotificationService + from .knowledge_propagation_engine import KnowledgePropagationEngine + from .organization_rag_manager import OrganizationRAGManager + from .task_knowledge_extractor import TaskKnowledgeExtractor + from .team_knowledge_manager import TeamKnowledgeManager + + app.state.org_rag_manager = OrganizationRAGManager(database_url) + await app.state.org_rag_manager.initialize() + + app.state.team_knowledge_manager = TeamKnowledgeManager(database_url) + await app.state.team_knowledge_manager.initialize() + + app.state.knowledge_propagation_engine = KnowledgePropagationEngine( + database_url, app.state.org_rag_manager, app.state.team_knowledge_manager + ) + await app.state.knowledge_propagation_engine.initialize() + + app.state.notification_service = KnowledgeNotificationService(database_url) + await app.state.notification_service.initialize() + + app.state.task_knowledge_extractor = TaskKnowledgeExtractor( + database_url, + app.state.org_rag_manager, + app.state.team_knowledge_manager, + app.state.knowledge_propagation_engine, + ) + await app.state.task_knowledge_extractor.initialize() + + app.state.context_enhancement_service = ContextEnhancementService( + database_url, app.state.org_rag_manager, app.state.team_knowledge_manager + ) + await app.state.context_enhancement_service.initialize() + + # Initialize knowledge analytics service + from .knowledge_analytics_service import KnowledgeAnalyticsService + + app.state.knowledge_analytics_service = KnowledgeAnalyticsService(database_url) + await app.state.knowledge_analytics_service.initialize() + + logger.info("Knowledge management system initialized successfully") + + except Exception as e: + logger.warning(f"Failed to initialize knowledge management system: {e}") + + # Initialize goals management system + try: + from .goal_conversation_service import GoalConversationService + from .goal_tracking_service import GoalTrackingService + from .goals_management_service import GoalsManagementService + from .milestone_task_engine import MilestoneTaskEngine + + app.state.goals_service = GoalsManagementService(database_url) + await app.state.goals_service.initialize() + + app.state.milestone_task_engine = MilestoneTaskEngine(database_url) + await app.state.milestone_task_engine.initialize() + + app.state.goal_conversation_service = GoalConversationService(database_url) + await app.state.goal_conversation_service.initialize() + + app.state.goal_tracking_service = GoalTrackingService(database_url) + await app.state.goal_tracking_service.initialize() + + logger.info("Goals management system initialized successfully") + + except Exception as e: + logger.warning(f"Failed to initialize goals management system: {e}") + + # Connect components + app.state.task_queue.set_task_execution_engine(app.state.task_execution_engine) + await app.state.agent_manager.set_sandbox_manager(app.state.sandbox_manager) + + # Initialize IzzyAI CEO on startup + try: + await app.state.agent_manager.create_agent( + name="IzzyAI", + role="Digital CEO", + type="executive", + config={ + "model": "claude-sonnet-4-20250514", + "temperature": 0.7, + "tools": [ + "strategic_planning", + "resource_allocation", + "team_management", + ], + }, + ) + except Exception as e: + print(f"Warning: Could not create IzzyAI CEO: {e}") + + yield + + # Shutdown + await app.state.multi_agent_coordinator.stop() + await app.state.task_execution_engine.stop() + await app.state.sandbox_manager.stop() + await app.state.agent_manager.shutdown_all() + await app.state.task_queue.close() + + # Shutdown knowledge management services + try: + if hasattr(app.state, "knowledge_analytics_service"): + await app.state.knowledge_analytics_service.close() + if hasattr(app.state, "context_enhancement_service"): + await app.state.context_enhancement_service.close() + if hasattr(app.state, "task_knowledge_extractor"): + await app.state.task_knowledge_extractor.close() + if hasattr(app.state, "notification_service"): + await app.state.notification_service.close() + if hasattr(app.state, "knowledge_propagation_engine"): + await app.state.knowledge_propagation_engine.close() + if hasattr(app.state, "team_knowledge_manager"): + await app.state.team_knowledge_manager.close() + if hasattr(app.state, "org_rag_manager"): + await app.state.org_rag_manager.close() + logger.info("Knowledge management system shutdown complete") + except Exception as e: + logger.error(f"Error shutting down knowledge management system: {e}") + + # Shutdown goals management services + try: + if hasattr(app.state, "goal_tracking_service"): + await app.state.goal_tracking_service.close() + if hasattr(app.state, "goal_conversation_service"): + await app.state.goal_conversation_service.close() + if hasattr(app.state, "milestone_task_engine"): + await app.state.milestone_task_engine.close() + if hasattr(app.state, "goals_service"): + await app.state.goals_service.close() + logger.info("Goals management system shutdown complete") + except Exception as e: + logger.error(f"Error shutting down goals management system: {e}") + + +app = FastAPI( + title="FuzeAgent Orchestrator API", + description=""" + ## FuzeAgent AI Team Orchestration Platform + + A comprehensive platform for autonomous AI development teams that enables: + + ### 🤖 Autonomous Agent Execution + - **Claude SDK Integration**: Interactive AI development with real-time conversation streaming + - **File Operations Engine**: Safe code changes with human approval workflows + - **Multi-Agent Coordination**: Complex task decomposition and agent collaboration + + ### 🏗️ Core Features + - **Agent Management**: Create, configure, and manage AI development agents + - **Task Orchestration**: Assign and monitor complex development tasks + - **Real-time Monitoring**: WebSocket streaming for live progress updates + - **Human-in-the-Loop**: Seamless approval workflows for critical decisions + + ### 🔗 Integration Capabilities + - **MCP (Model Context Protocol)**: Organizational context for AI agents + - **Git Workflow Management**: Automated repository operations + - **Sandbox Environments**: Isolated development containers + - **Database Integration**: PostgreSQL for persistent storage + + ### 📡 API Categories + - **Agent Management**: Create and manage AI agents + - **Task Execution**: Autonomous task processing + - **File Operations**: Code change management + - **Multi-Agent Coordination**: Team collaboration + - **Real-time Communication**: WebSocket endpoints + - **MCP Integration**: Organizational context tools + - **Goals Management**: Organizational goals, milestones, and task planning + - **Knowledge Management**: RAG system and organizational memory + + **Version**: 2.0.0 (Autonomous Execution) + """, + version="2.0.0", + lifespan=lifespan, + docs_url="/docs", + redoc_url="/redoc", + openapi_tags=[ + {"name": "health", "description": "Health check and system status endpoints"}, + { + "name": "agents", + "description": "AI agent creation, management, and status monitoring", + }, + {"name": "tasks", "description": "Task assignment, execution, and monitoring"}, + { + "name": "autonomous-execution", + "description": "Autonomous task execution with Claude SDK integration", + }, + { + "name": "file-operations", + "description": "File system operations and code change management", + }, + { + "name": "multi-agent-coordination", + "description": "Multi-agent collaboration and task coordination", + }, + { + "name": "real-time", + "description": "WebSocket endpoints for real-time updates", + }, + { + "name": "human-in-loop", + "description": "Human approval workflows and interaction handling", + }, + { + "name": "mcp-integration", + "description": "Model Context Protocol tools and resources", + }, + {"name": "sandboxes", "description": "Sandbox environment management"}, + {"name": "context", "description": "Agent memory and context management"}, + { + "name": "model-configuration", + "description": "AI model configuration and API key management", + }, + { + "name": "knowledge-management", + "description": "Hierarchical knowledge management, RAG, and intelligent notifications", + }, + ], +) + +app.add_middleware( + CORSMiddleware, + allow_origins=[ + "http://localhost:3000", + "http://localhost:3031", + "http://localhost", + ], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +# Include hierarchy router for organizational visualization +app.include_router(hierarchy_router) + + +# Health check endpoint +@app.get( + "/health", + tags=["health"], + summary="Health Check", + description="Check the health status of the FuzeAgent orchestrator service", + response_description="Service health status", +) +async def health_check(): + """ + Health check endpoint that returns the current status of the orchestrator service. + + Returns: + dict: Service health status and basic information + """ + return { + "status": "healthy", + "service": "orchestrator", + "version": "2.0.0", + "features": { + "autonomous_execution": True, + "multi_agent_coordination": True, + "file_operations": True, + "mcp_integration": True, + "real_time_streaming": True, + }, + # Whether GET /openapi.yaml can answer. An image built without its + # contract is DEGRADED, not dead — the probe still passes (no restart + # can conjure a file the image lacks) but the condition is visible to + # anything that looks, instead of surfacing only as a 503 later. + "openapi": "loaded" if _openapi_document() is not None else "unavailable", + } + + +# --------------------------------------------------------------------------- +# The contract, SERVED. +# +# contracts/openapi.yaml describes this orchestrator's real HTTP surface, with +# the curated descriptions and the irreversibility guidance that +# mcp/tools.overrides.yaml narrows. Committing it is not the same as publishing +# it: consumers — the MCP gateway among them — discover the surface over HTTP. +# +# This is NOT /openapi.json. FastAPI generates that from the code at import +# time; it is accurate about shapes and says nothing about which operations +# dispatch an agent that cannot be recalled. Both are served. This one is the +# contract. +# +# The document is read from the IMAGE, never from a mount, so what this endpoint +# publishes is always the contract this build was compiled against. +# --------------------------------------------------------------------------- +_ORCH_DIR = os.path.dirname(os.path.abspath(__file__)) +_OPENAPI_CANDIDATES = [ + p + for p in [ + os.getenv("OPENAPI_SPEC_PATH"), + os.path.join(_ORCH_DIR, "contracts", "openapi.yaml"), + os.path.join(_ORCH_DIR, "..", "..", "contracts", "openapi.yaml"), + ] + if p +] + + +def _openapi_document(): + """Return the OpenAPI document text, or None when the image lacks it.""" + for path in _OPENAPI_CANDIDATES: + try: + with open(path, "r", encoding="utf-8") as fh: + return fh.read() + except OSError: + continue + return None + + +@app.get( + "/openapi.yaml", + tags=["health"], + summary="This OpenAPI Document", + description=( + "Serve contracts/openapi.yaml — the curated contract, as distinct from " + "FastAPI's auto-generated /openapi.json." + ), + include_in_schema=False, +) +async def get_openapi_document(): + doc = _openapi_document() + if doc is None: + logger.error("OpenAPI document not found; tried %s", _OPENAPI_CANDIDATES) + # 503, not 500 and not a crash: the service is otherwise functional and + # no restart can produce a spec the image does not contain. + raise HTTPException( + status_code=503, + detail=( + "openapi_document_unavailable: this image was built without " + "contracts/openapi.yaml. Rebuild with the repo root as the Docker " + "context so the contract is copied in." + ), + ) + return Response(content=doc, media_type="application/yaml") + + +# WebSocket for real-time updates +@app.websocket("/ws") +async def websocket_endpoint(websocket: WebSocket): + await websocket.accept() + try: + while True: + # Send agent updates to UI + updates = await app.state.agent_manager.get_updates() + await websocket.send_json(updates) + await asyncio.sleep(1) + except Exception as e: + print(f"WebSocket error: {e}") + finally: + await websocket.close() + + +# WebSocket for task execution updates +@app.websocket("/ws/tasks/{task_id}") +async def task_websocket_endpoint(websocket: WebSocket, task_id: str): + """WebSocket endpoint for real-time task execution updates""" + await websocket.accept() + try: + while True: + # Get task execution status + try: + status = await app.state.task_queue.get_execution_status(task_id) + await websocket.send_json( + {"type": "status_update", "task_id": task_id, "data": status} + ) + + # If task is completed or failed, send final update and close + if status.get("status") in ["completed", "failed", "cancelled"]: + await websocket.send_json( + { + "type": "task_finished", + "task_id": task_id, + "final_status": status.get("status"), + } + ) + break + + except Exception as e: + await websocket.send_json( + {"type": "error", "task_id": task_id, "error": str(e)} + ) + + await asyncio.sleep(2) # Update every 2 seconds + + except Exception as e: + print(f"Task WebSocket error for {task_id}: {e}") + finally: + await websocket.close() + + +# WebSocket for real-time Claude SDK conversation streaming +@app.websocket("/ws/tasks/{task_id}/conversation") +async def conversation_websocket_endpoint(websocket: WebSocket, task_id: str): + """WebSocket endpoint for real-time Claude SDK conversation streaming""" + await websocket.accept() + try: + # Get execution context + execution = app.state.task_execution_engine.active_executions.get(task_id) + if not execution: + await websocket.send_json( + {"type": "error", "message": f"Task {task_id} not found or not active"} + ) + await websocket.close() + return + + # Wait for Claude SDK session to be available + while not execution.claude_session_id and execution.status not in [ + "completed", + "failed", + "cancelled", + ]: + await asyncio.sleep(1) + + if not execution.claude_session_id: + await websocket.send_json( + { + "type": "error", + "message": "No active Claude SDK session for this task", + } + ) + await websocket.close() + return + + # Stream Claude SDK output + claude_sdk_manager = execution.claude_sdk_manager + if claude_sdk_manager: + try: + async for output_chunk in claude_sdk_manager.stream_output( + execution.claude_session_id + ): + await websocket.send_json( + { + "type": "claude_output", + "task_id": task_id, + "content": output_chunk, + "timestamp": datetime.now().isoformat(), + } + ) + + # Session ended + await websocket.send_json( + { + "type": "conversation_ended", + "task_id": task_id, + "timestamp": datetime.now().isoformat(), + } + ) + + except Exception as e: + await websocket.send_json( + { + "type": "error", + "message": f"Error streaming conversation: {str(e)}", + } + ) + + except Exception as e: + print(f"Conversation WebSocket error for {task_id}: {e}") + finally: + await websocket.close() + + +# WebSocket for file operations streaming +@app.websocket("/ws/tasks/{task_id}/file-operations") +async def file_operations_websocket_endpoint(websocket: WebSocket, task_id: str): + """WebSocket endpoint for real-time file operations updates""" + await websocket.accept() + try: + # Get execution context + execution = app.state.task_execution_engine.active_executions.get(task_id) + if not execution: + await websocket.send_json( + {"type": "error", "message": f"Task {task_id} not found or not active"} + ) + await websocket.close() + return + + file_ops_engine = execution.file_operations_engine + if not file_ops_engine: + await websocket.send_json( + { + "type": "error", + "message": "No file operations engine available for this task", + } + ) + await websocket.close() + return + + last_batch_count = 0 + + while execution.status not in ["completed", "failed", "cancelled"]: + try: + # Get pending operations + pending_operations = file_ops_engine.get_pending_operations() + applied_operations = file_ops_engine.get_applied_operations() + + current_batch_count = len(pending_operations) + len(applied_operations) + + # Send updates if there are new operations + if current_batch_count > last_batch_count: + # Send pending operations + for batch in pending_operations: + # Get diff preview + diffs = await file_ops_engine.get_file_diff_preview( + batch.batch_id + ) + + await websocket.send_json( + { + "type": "pending_operations", + "task_id": task_id, + "batch_id": batch.batch_id, + "description": batch.description, + "requires_approval": batch.requires_approval, + "operations_count": len(batch.operations), + "file_diffs": diffs, + "timestamp": batch.created_at.isoformat(), + } + ) + + # Send applied operations + for batch in applied_operations: + await websocket.send_json( + { + "type": "applied_operations", + "task_id": task_id, + "batch_id": batch.batch_id, + "description": batch.description, + "operations_count": len(batch.operations), + "applied_at": ( + batch.applied_at.isoformat() + if batch.applied_at + else None + ), + "timestamp": batch.created_at.isoformat(), + } + ) + + last_batch_count = current_batch_count + + await asyncio.sleep(1) # Check every second + + except Exception as e: + await websocket.send_json( + { + "type": "error", + "message": f"Error getting file operations: {str(e)}", + } + ) + + # Task completed + await websocket.send_json( + { + "type": "task_completed", + "task_id": task_id, + "final_status": execution.status.value, + "timestamp": datetime.now().isoformat(), + } + ) + + except Exception as e: + print(f"File operations WebSocket error for {task_id}: {e}") + finally: + await websocket.close() + + +# Agent Management Endpoints +@app.post( + "/agents", + tags=["agents"], + summary="Create AI Agent", + description="Create a new AI agent with repository and sandbox settings", + response_model=AgentResponse, +) +async def create_agent(agent_config: AgentCreateRequest, _auth=Depends(require_auth)): + """Create a new AI agent with repository and sandbox settings""" + try: + agent = await app.state.agent_manager.create_agent(**agent_config) + return { + "agent_id": agent.id, + "status": "created", + "agent": { + "id": agent.id, + "name": agent_config.get("name"), + "role": agent_config.get("role"), + "type": agent_config.get("type"), + "repository_settings": agent_config.get("repository_settings", {}), + "sandbox_settings": agent_config.get("sandbox_settings", {}), + "created_at": ( + agent.created_at if hasattr(agent, "created_at") else None + ), + }, + } + except Exception as e: + raise HTTPException(status_code=400, detail=f"Failed to create agent: {str(e)}") + + +@app.get( + "/agents", + tags=["agents"], + summary="List All Agents", + description="Get a list of all AI agents and their current status", +) +async def list_agents(): + """List all agents and their status""" + return await app.state.agent_manager.list_agents() + + +@app.post( + "/agents/{agent_id}/tasks", + tags=["tasks"], + summary="Assign Task to Agent", + description="Assign a specific task to an AI agent", + response_model=TaskResponse, +) +async def assign_task( + agent_id: str = Path(..., description="Agent ID"), + task: TaskCreateRequest = Body(...), + _auth=Depends(require_auth), +): + """Assign a task to an agent""" + task_id = await app.state.task_queue.assign_task(agent_id, task) + return {"task_id": task_id, "status": "assigned"} + + +@app.get("/agents/{agent_id}/status") +async def get_agent_status(agent_id: str): + """Get detailed agent status""" + return await app.state.agent_manager.get_agent_status(agent_id) + + +@app.get("/agents/{agent_id}/tasks") +async def get_agent_tasks(agent_id: str): + """Get tasks assigned to an agent""" + try: + # This would normally query the database for tasks assigned to the agent + # For now, return mock data + return [ + { + "id": "1", + "title": "Strategic Planning Q4 2025", + "description": "Develop comprehensive strategic plan for Q4 2025 expansion", + "status": "completed", + "priority": "high", + "created_at": "2025-08-05T09:00:00Z", + "completed_at": "2025-08-05T17:30:00Z", + }, + { + "id": "2", + "title": "Team Performance Review", + "description": "Conduct quarterly performance review for all team leads", + "status": "in_progress", + "priority": "medium", + "created_at": "2025-08-06T08:00:00Z", + }, + ] + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get agent tasks: {str(e)}" + ) + + +@app.get("/teams") +async def list_teams(): + """List all teams""" + try: + # This would normally query the database for teams + # For now, return mock data + return [ + { + "id": "1", + "name": "Executive Team", + "description": "Strategic leadership and decision making", + }, + { + "id": "2", + "name": "Development Team", + "description": "Frontend, backend, and full-stack development", + }, + { + "id": "3", + "name": "Quality Assurance", + "description": "Testing, quality control, and bug detection", + }, + { + "id": "4", + "name": "DevOps Team", + "description": "Infrastructure, deployment, and system operations", + }, + { + "id": "5", + "name": "Business Team", + "description": "Marketing, sales, and customer relations", + }, + ] + except Exception as e: + raise HTTPException(status_code=500, detail=f"Failed to list teams: {str(e)}") + + +@app.get("/agent-templates") +async def list_agent_templates(): + """List available agent templates""" + try: + return [ + { + "id": "react_developer", + "name": "React Developer", + "description": "Frontend developer specialized in React and TypeScript", + "type": "developer", + "defaultConfig": { + "model": "claude-sonnet-4-20250514", + "temperature": 0.7, + "tools": ["code_generation", "code_review", "debugging", "testing"], + "goal": "Build responsive and performant React applications", + "backstory": "Experienced frontend developer with deep knowledge of React ecosystem", + }, + }, + { + "id": "python_developer", + "name": "Python Developer", + "description": "Backend developer specialized in Python and FastAPI", + "type": "developer", + "defaultConfig": { + "model": "claude-sonnet-4-20250514", + "temperature": 0.7, + "tools": [ + "code_generation", + "api_development", + "database_design", + "testing", + ], + "goal": "Develop robust and scalable backend systems", + "backstory": "Senior Python developer with expertise in FastAPI and databases", + }, + }, + { + "id": "qa_engineer", + "name": "QA Engineer", + "description": "Quality assurance engineer focused on testing and automation", + "type": "qa", + "defaultConfig": { + "model": "claude-sonnet-4-20250514", + "temperature": 0.6, + "tools": [ + "test_automation", + "bug_reporting", + "quality_analysis", + "performance_testing", + ], + "goal": "Ensure high quality and reliability of software products", + "backstory": "Experienced QA engineer with expertise in automated testing frameworks", + }, + }, + { + "id": "devops_engineer", + "name": "DevOps Engineer", + "description": "Infrastructure and deployment specialist", + "type": "devops", + "defaultConfig": { + "model": "claude-sonnet-4-20250514", + "temperature": 0.5, + "tools": [ + "infrastructure_management", + "deployment", + "monitoring", + "security", + ], + "goal": "Maintain reliable and scalable infrastructure", + "backstory": "DevOps engineer with expertise in cloud platforms and CI/CD", + }, + }, + ] + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to list agent templates: {str(e)}" + ) + + +@app.get("/tasks") +async def list_tasks(): + """List all tasks""" + return await app.state.task_queue.list_tasks() + + +@app.get("/tasks/{task_id}") +async def get_task(task_id: str): + """Get task details""" + return await app.state.task_queue.get_task(task_id) + + +# Autonomous Execution Endpoints +@app.post("/agents/from-template") +async def create_agent_from_template(request: dict): + """Create agent from template with repository settings""" + try: + # Extract template data + template_id = request.get("template_id") + name = request.get("name") + team_id = request.get("team_id") + overrides = request.get("overrides", {}) + + # Get template configuration + template_config = await app.state.agent_manager.get_template_config(template_id) + if not template_config: + raise HTTPException( + status_code=404, detail=f"Template {template_id} not found" + ) + + # Build agent configuration + agent_config = { + "name": name, + "role": template_config.get("role", template_id.replace("_", " ").title()), + "type": template_config.get("type", "specialized"), + "template_id": template_id, + "team_id": team_id, + "config": {**template_config.get("config", {}), **overrides}, + "repository_settings": request.get("repository_settings", {}), + "sandbox_settings": { + "base_image": f"fuzeagent/dev-{template_id.split('_')[0]}:latest", + "resource_limits": template_config.get( + "resource_limits", {"memory": "2Gi", "cpu": "1.0", "disk": "10Gi"} + ), + "auto_cleanup": "24h", + }, + } + + # Create agent + agent = await app.state.agent_manager.create_agent(**agent_config) + + return { + "agent_id": agent.id, + "status": "created", + "agent": agent_config, + "template_id": template_id, + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to create agent from template: {str(e)}" + ) + + +@app.get("/templates") +async def get_agent_templates(): + """Get available agent templates""" + return await app.state.agent_manager.get_available_templates() + + +@app.post( + "/tasks/{task_id}/execute", + tags=["autonomous-execution"], + summary="Start Autonomous Task Execution", + description="Begin autonomous execution of a task using Claude SDK integration", + response_model=TaskResponse, +) +async def start_task_execution( + task_id: str = Path(..., description="Task ID to execute") +): + """Start autonomous execution of a task""" + try: + # This will be handled by the TaskExecutionEngine + result = await app.state.task_queue.start_autonomous_execution(task_id) + return {"task_id": task_id, "status": "execution_started", "result": result} + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to start task execution: {str(e)}" + ) + + +@app.get("/tasks/{task_id}/status") +async def get_task_execution_status(task_id: str): + """Get detailed task execution status""" + try: + status = await app.state.task_queue.get_execution_status(task_id) + return status + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get task status: {str(e)}" + ) + + +@app.get("/tasks/{task_id}/iterations") +async def get_task_iterations(task_id: str): + """Get task iteration history""" + try: + iterations = await app.state.task_queue.get_task_iterations(task_id) + return {"task_id": task_id, "iterations": iterations} + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get task iterations: {str(e)}" + ) + + +@app.get("/agents/{agent_id}/sandbox") +async def get_agent_sandbox(agent_id: str): + """Get agent sandbox information""" + try: + sandbox_info = await app.state.agent_manager.get_agent_sandbox(agent_id) + return sandbox_info + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get agent sandbox: {str(e)}" + ) + + +# Additional endpoints for UI support +@app.put("/tasks/{task_id}") +async def update_task(task_id: str, update_data: dict): + """Update task status and result""" + await app.state.task_queue.update_task_status( + task_id=task_id, + status=update_data.get("status"), + result=update_data.get("result"), + ) + return {"status": "updated"} + + +@app.post("/context/interactions") +async def store_interaction(interaction_data: dict): + """Store agent interaction""" + interaction_id = await app.state.context_service.store_interaction( + agent_id=interaction_data.get("agent_id"), + content=interaction_data.get("content"), + metadata=interaction_data.get("metadata", {}), + ) + return {"interaction_id": interaction_id} + + +@app.get("/context") +async def get_context(query: str, agent_id: str = None): + """Get relevant context for a query""" + context = await app.state.context_service.get_context(query, agent_id) + return context + + +@app.get("/agents/{agent_id}/memory") +async def get_agent_memory(agent_id: str, limit: int = 10): + """Get agent memory""" + memory = await app.state.context_service.get_agent_memory(agent_id, limit) + return memory + + +# Agent Conversation Endpoints +@app.get( + "/agents/{agent_id}/conversations", + tags=["conversations"], + summary="Get Agent Conversations", + description="Get all conversations for a specific agent", +) +async def get_agent_conversations(agent_id: str): + """Get all conversations for a specific agent""" + try: + async with get_db_connection() as conn: + conversations = await conn.fetch( + """ + SELECT cs.*, COUNT(ac.id) as message_count, + (SELECT content FROM agent_conversations + WHERE session_id = cs.id + ORDER BY created_at DESC LIMIT 1) as last_message + FROM chat_sessions cs + LEFT JOIN agent_conversations ac ON cs.id = ac.session_id + WHERE cs.agent_id = $1 + GROUP BY cs.id + ORDER BY cs.last_activity DESC + """, + agent_id, + ) + + return [dict(row) for row in conversations] + + except Exception as e: + logger.error(f"Error getting agent conversations: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/agents/{agent_id}/conversations", + tags=["conversations"], + summary="Create New Agent Conversation", + description="Create a new conversation with an agent", +) +async def create_agent_conversation(agent_id: str, request: ConversationCreateRequest): + """Create a new conversation with an agent""" + try: + async with get_db_connection() as conn: + # Create new chat session + session_id = await conn.fetchval( + """ + INSERT INTO chat_sessions (agent_id, session_name, context, status) + VALUES ($1, $2, $3, 'active') + RETURNING id + """, + agent_id, + request.title, + request.context or {}, + ) + + # Add initial message if provided + if request.initial_message: + await conn.execute( + """ + INSERT INTO agent_conversations (session_id, agent_id, message_type, content) + VALUES ($1, $2, 'system', $3) + """, + session_id, + agent_id, + request.initial_message, + ) + + # Get the created conversation + conversation = await conn.fetchrow( + """ + SELECT * FROM chat_sessions WHERE id = $1 + """, + session_id, + ) + + return dict(conversation) + + except Exception as e: + logger.error(f"Error creating agent conversation: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/agents/{agent_id}/conversations/{conversation_id}/messages", + tags=["conversations"], + summary="Get Conversation Messages", + description="Get all messages in a conversation", +) +async def get_conversation_messages(agent_id: str, conversation_id: str): + """Get all messages in a conversation""" + try: + async with get_db_connection() as conn: + messages = await conn.fetch( + """ + SELECT * FROM agent_conversations + WHERE session_id = $1 AND agent_id = $2 + ORDER BY created_at ASC + """, + conversation_id, + agent_id, + ) + + return [dict(row) for row in messages] + + except Exception as e: + logger.error(f"Error getting conversation messages: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/agents/{agent_id}/conversations/{conversation_id}/messages", + tags=["conversations"], + summary="Send Message to Agent", + description="Send a message to an agent in a conversation", +) +async def send_message_to_agent( + agent_id: str, conversation_id: str, request: ChatMessageRequest +): + """Send a message to an agent in a conversation""" + try: + async with get_db_connection() as conn: + # Insert user message + user_message_id = await conn.fetchval( + """ + INSERT INTO agent_conversations (session_id, agent_id, message_type, content, metadata) + VALUES ($1, $2, 'user', $3, $4) + RETURNING id + """, + conversation_id, + agent_id, + request.content, + request.metadata or {}, + ) + + # Update session last activity + await conn.execute( + """ + UPDATE chat_sessions + SET last_activity = CURRENT_TIMESTAMP, message_count = message_count + 1 + WHERE id = $1 + """, + conversation_id, + ) + + # TODO: Here we would integrate with the actual agent to generate a response + # For now, return a simple acknowledgment + + return { + "id": str(user_message_id), + "status": "sent", + "message": "Message sent to agent", + } + + except Exception as e: + logger.error(f"Error sending message to agent: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get("/knowledge/search") +async def search_knowledge(query: str, limit: int = 10): + """Search knowledge base""" + results = await app.state.context_service.search_knowledge(query, limit) + return results + + +# Human-in-the-loop endpoints +@app.post( + "/tasks/{task_id}/human-response", + tags=["human-in-loop"], + summary="Submit Human Response", + description="Submit human response to a task question or approval request", +) +async def submit_human_response( + task_id: str = Path(..., description="Task ID"), + response_data: HumanResponseRequest = Body(...), +): + """Submit human response to a task question""" + try: + response = response_data.get("response", "") + if not response: + raise HTTPException(status_code=400, detail="Response cannot be empty") + + success = await app.state.task_queue.handle_human_response(task_id, response) + + if success: + return {"status": "success", "message": "Human response submitted"} + else: + raise HTTPException( + status_code=404, + detail="Task not found or not waiting for human response", + ) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to submit human response: {str(e)}" + ) + + +@app.post("/tasks/{task_id}/cancel") +async def cancel_task_execution(task_id: str): + """Cancel autonomous execution of a task""" + try: + success = await app.state.task_queue.cancel_task_execution(task_id) + + if success: + return {"status": "cancelled", "message": "Task execution cancelled"} + else: + raise HTTPException(status_code=404, detail="Task not found or not running") + + except Exception as e: + raise HTTPException(status_code=500, detail=f"Failed to cancel task: {str(e)}") + + +@app.get("/tasks/{task_id}/messages") +async def get_task_messages(task_id: str): + """Get task messages and chat history""" + try: + # This would integrate with the HumanInTheLoopHandler when implemented + # For now, return iteration history which includes human interactions + iterations = await app.state.task_queue.get_task_iterations(task_id) + + messages = [] + for iteration in iterations: + if iteration.get("human_question"): + messages.append( + { + "type": "agent_question", + "content": iteration["human_question"], + "timestamp": iteration["started_at"], + "iteration": iteration["iteration_number"], + } + ) + + if iteration.get("human_response"): + messages.append( + { + "type": "human_response", + "content": iteration["human_response"], + "timestamp": iteration["completed_at"] + or iteration["started_at"], + "iteration": iteration["iteration_number"], + } + ) + + return {"task_id": task_id, "messages": messages} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get task messages: {str(e)}" + ) + + +# Sandbox management endpoints +@app.get("/sandboxes") +async def list_sandboxes(agent_id: str = None, status: str = None): + """List active sandboxes""" + try: + from .sandbox_manager import SandboxStatus + + sandbox_status = None + if status: + try: + sandbox_status = SandboxStatus(status) + except ValueError: + raise HTTPException(status_code=400, detail=f"Invalid status: {status}") + + sandboxes = await app.state.sandbox_manager.list_sandboxes( + agent_id=agent_id, status=sandbox_status + ) + + return { + "sandboxes": [ + { + "sandbox_id": s.sandbox_id, + "agent_id": s.agent_id, + "task_id": s.task_id, + "status": s.status.value, + "workspace_path": s.workspace_path, + "created_at": s.created_at.isoformat(), + "resource_limits": s.resource_limits, + } + for s in sandboxes + ] + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to list sandboxes: {str(e)}" + ) + + +@app.post("/sandboxes/{sandbox_id}/execute") +async def execute_command_in_sandbox(sandbox_id: str, command_data: dict): + """Execute a command in a sandbox""" + try: + command = command_data.get("command") + working_dir = command_data.get("working_dir") + + if not command: + raise HTTPException(status_code=400, detail="Command is required") + + result = await app.state.sandbox_manager.execute_command( + sandbox_id=sandbox_id, command=command, working_dir=working_dir + ) + + return result + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to execute command: {str(e)}" + ) + + +@app.delete("/sandboxes/{sandbox_id}") +async def destroy_sandbox(sandbox_id: str): + """Destroy a sandbox""" + try: + await app.state.sandbox_manager.destroy_sandbox(sandbox_id) + return {"status": "destroyed", "sandbox_id": sandbox_id} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to destroy sandbox: {str(e)}" + ) + + +# Agent registration and communication endpoints +@app.post("/agents/{agent_id}/register") +async def register_agent(agent_id: str, registration_data: dict): + """Register an agent running in a sandbox container""" + try: + # Store agent registration info + # This would typically update the agent's status and capabilities + return { + "status": "registered", + "agent_id": agent_id, + "registered_at": datetime.now().isoformat(), + } + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to register agent: {str(e)}" + ) + + +@app.get("/agents/{agent_id}/next-task") +async def get_next_task_for_agent(agent_id: str): + """Get the next task for an agent to execute""" + try: + # Find pending tasks assigned to this agent + tasks = await app.state.task_queue.get_agent_tasks(agent_id) + pending_tasks = [t for t in tasks if t.get("status") == "pending"] + + if pending_tasks: + # Return the first pending task + task = pending_tasks[0] + # Update status to 'assigned' to prevent double assignment + await app.state.task_queue.update_task_status(task["id"], "assigned") + return task + else: + # No tasks available + return None, 204 + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get next task: {str(e)}" + ) + + +@app.post("/agents/{agent_id}/error") +async def report_agent_error(agent_id: str, error_data: dict): + """Report an error from an agent""" + try: + # Log the error and update agent status + logger.error(f"Agent {agent_id} reported error: {error_data.get('error')}") + + # You might want to store this in a database or alerting system + return {"status": "error_logged", "agent_id": agent_id} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to log agent error: {str(e)}" + ) + + +# Conversation management endpoints +@app.get("/tasks/{task_id}/conversation") +async def get_task_conversation(task_id: str, iteration: int = None, limit: int = 100): + """Get conversation history for a task""" + try: + conversation_history = await app.state.task_execution_engine.conversation_manager.get_conversation_history( + task_id=task_id, iteration_number=iteration, limit=limit + ) + + return {"task_id": task_id, "conversation": conversation_history} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get conversation: {str(e)}" + ) + + +@app.get("/tasks/{task_id}/conversation/summary") +async def get_conversation_summary(task_id: str): + """Get conversation summary with statistics""" + try: + summary = await app.state.task_execution_engine.conversation_manager.get_conversation_summary( + task_id + ) + return summary + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get conversation summary: {str(e)}" + ) + + +@app.get("/tasks/{task_id}/code-generations") +async def get_task_code_generations( + task_id: str, iteration: int = None, file_type: str = None +): + """Get code generations for a task""" + try: + code_generations = await app.state.task_execution_engine.conversation_manager.get_code_generations( + task_id=task_id, iteration_number=iteration, file_type=file_type + ) + + return {"task_id": task_id, "code_generations": code_generations} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get code generations: {str(e)}" + ) + + +@app.get("/agents/{agent_id}/performance") +async def get_agent_performance(agent_id: str, hours: int = 24): + """Get agent performance metrics""" + try: + metrics = await app.state.task_execution_engine.conversation_manager.get_agent_performance_metrics( + agent_id=agent_id, time_range_hours=hours + ) + + return {"agent_id": agent_id, "time_range_hours": hours, "metrics": metrics} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get agent performance: {str(e)}" + ) + + +# File Operations Endpoints +@app.get( + "/tasks/{task_id}/file-operations", + tags=["file-operations"], + summary="Get File Operations", + description="Get file operations for a task with optional status filtering", +) +async def get_task_file_operations( + task_id: str = Path(..., description="Task ID"), + status: Optional[str] = Query( + None, description="Filter by status (pending, applied)" + ), +): + """Get file operations for a task""" + try: + execution = app.state.task_execution_engine.active_executions.get(task_id) + if not execution or not execution.file_operations_engine: + raise HTTPException( + status_code=404, detail="Task not found or no file operations available" + ) + + file_ops_engine = execution.file_operations_engine + + if status == "pending": + operations = file_ops_engine.get_pending_operations() + elif status == "applied": + operations = file_ops_engine.get_applied_operations() + else: + # Get all operations + pending = file_ops_engine.get_pending_operations() + applied = file_ops_engine.get_applied_operations() + operations = pending + applied + + # Convert to dict format + operations_data = [] + for batch in operations: + operations_data.append( + { + "batch_id": batch.batch_id, + "task_id": batch.task_id, + "agent_id": batch.agent_id, + "description": batch.description, + "requires_approval": batch.requires_approval, + "approval_status": batch.approval_status.value, + "operations_count": len(batch.operations), + "created_at": batch.created_at.isoformat(), + "applied_at": ( + batch.applied_at.isoformat() if batch.applied_at else None + ), + } + ) + + return {"task_id": task_id, "operations": operations_data} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get file operations: {str(e)}" + ) + + +@app.get("/tasks/{task_id}/file-operations/{batch_id}/preview") +async def get_file_operations_preview(task_id: str, batch_id: str): + """Get preview of file changes for a batch""" + try: + execution = app.state.task_execution_engine.active_executions.get(task_id) + if not execution or not execution.file_operations_engine: + raise HTTPException( + status_code=404, detail="Task not found or no file operations available" + ) + + file_ops_engine = execution.file_operations_engine + diffs = await file_ops_engine.get_file_diff_preview(batch_id) + + return {"task_id": task_id, "batch_id": batch_id, "file_diffs": diffs} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get file preview: {str(e)}" + ) + + +@app.post( + "/tasks/{task_id}/file-operations/{batch_id}/approve", + tags=["file-operations", "human-in-loop"], + summary="Approve File Operations", + description="Approve or reject file operations from Claude SDK", +) +async def approve_file_operations( + task_id: str = Path(..., description="Task ID"), + batch_id: str = Path(..., description="Batch ID"), + approval_data: FileOperationApprovalRequest = Body(...), +): + """Approve or reject file operations""" + try: + approved = approval_data.get("approved", False) + + execution = app.state.task_execution_engine.active_executions.get(task_id) + if not execution or not execution.file_operations_engine: + raise HTTPException( + status_code=404, detail="Task not found or no file operations available" + ) + + file_ops_engine = execution.file_operations_engine + success = await file_ops_engine.approve_operations(batch_id, approved) + + if success: + # Also notify Claude SDK if there's an active session + if execution.claude_sdk_manager and execution.claude_session_id: + await execution.claude_sdk_manager.approve_file_operations( + execution.claude_session_id, batch_id, approved + ) + + return { + "task_id": task_id, + "batch_id": batch_id, + "approved": approved, + "status": "success", + } + else: + raise HTTPException(status_code=400, detail="Failed to process approval") + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to approve file operations: {str(e)}" + ) + + +@app.post("/tasks/{task_id}/file-operations/{batch_id}/rollback") +async def rollback_file_operations(task_id: str, batch_id: str): + """Rollback applied file operations""" + try: + execution = app.state.task_execution_engine.active_executions.get(task_id) + if not execution or not execution.file_operations_engine: + raise HTTPException( + status_code=404, detail="Task not found or no file operations available" + ) + + file_ops_engine = execution.file_operations_engine + success = await file_ops_engine.rollback_operations(batch_id) + + if success: + return {"task_id": task_id, "batch_id": batch_id, "status": "rolled_back"} + else: + raise HTTPException(status_code=400, detail="Failed to rollback operations") + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to rollback file operations: {str(e)}" + ) + + +# Claude SDK Session Management Endpoints +@app.get("/tasks/{task_id}/claude-session") +async def get_claude_session_status(task_id: str): + """Get Claude SDK session status for a task""" + try: + execution = app.state.task_execution_engine.active_executions.get(task_id) + if ( + not execution + or not execution.claude_sdk_manager + or not execution.claude_session_id + ): + raise HTTPException( + status_code=404, detail="No active Claude SDK session for this task" + ) + + status = await execution.claude_sdk_manager.get_session_status( + execution.claude_session_id + ) + return status + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get Claude session status: {str(e)}" + ) + + +@app.post("/tasks/{task_id}/claude-session/input") +async def send_claude_session_input(task_id: str, input_data: dict): + """Send input to Claude SDK session""" + try: + user_input = input_data.get("input", "") + if not user_input: + raise HTTPException(status_code=400, detail="Input cannot be empty") + + execution = app.state.task_execution_engine.active_executions.get(task_id) + if ( + not execution + or not execution.claude_sdk_manager + or not execution.claude_session_id + ): + raise HTTPException( + status_code=404, detail="No active Claude SDK session for this task" + ) + + success = await execution.claude_sdk_manager.send_input( + execution.claude_session_id, user_input + ) + + if success: + return {"task_id": task_id, "status": "input_sent", "input": user_input} + else: + raise HTTPException( + status_code=400, detail="Failed to send input to Claude session" + ) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to send Claude session input: {str(e)}" + ) + + +# MCP Integration Endpoints +@app.get("/mcp/tools") +async def get_mcp_tools(): + """Get available MCP tools""" + try: + from .mcp_integration import FuzeAgentMCPServer + + mcp_server = FuzeAgentMCPServer() + tools = [ + { + "name": tool.name, + "description": tool.description, + "input_schema": tool.input_schema, + } + for tool in mcp_server.tools + ] + + return {"tools": tools} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get MCP tools: {str(e)}" + ) + + +@app.post( + "/mcp/call-tool", + tags=["mcp-integration"], + summary="Call MCP Tool", + description="Execute an MCP tool to access organizational context", +) +async def call_mcp_tool(tool_request: MCPToolRequest = Body(...)): + """Call an MCP tool""" + try: + from .mcp_integration import FuzeAgentMCPServer + + tool_name = tool_request.get("tool_name") + arguments = tool_request.get("arguments", {}) + + if not tool_name: + raise HTTPException(status_code=400, detail="tool_name is required") + + mcp_server = FuzeAgentMCPServer() + result = await mcp_server.handle_tool_call(tool_name, arguments) + + return result + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to call MCP tool: {str(e)}" + ) + + +@app.get("/mcp/resources") +async def get_mcp_resources(): + """Get available MCP resources""" + try: + from .mcp_integration import FuzeAgentMCPServer + + mcp_server = FuzeAgentMCPServer() + resources = [ + { + "uri": resource.uri, + "name": resource.name, + "description": resource.description, + "mime_type": resource.mime_type, + } + for resource in mcp_server.resources + ] + + return {"resources": resources} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get MCP resources: {str(e)}" + ) + + +@app.get("/mcp/resource") +async def get_mcp_resource(uri: str): + """Get an MCP resource by URI""" + try: + from .mcp_integration import FuzeAgentMCPServer + + if not uri: + raise HTTPException(status_code=400, detail="uri parameter is required") + + mcp_server = FuzeAgentMCPServer() + resource = await mcp_server.handle_resource_request(uri) + + return resource + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get MCP resource: {str(e)}" + ) + + +@app.get("/tasks/{task_id}/mcp-context") +async def get_task_mcp_context(task_id: str): + """Get MCP context for a task""" + try: + from .mcp_integration import FuzeAgentMCPServer, MCPClaudeIntegration + + execution = app.state.task_execution_engine.active_executions.get(task_id) + if not execution: + raise HTTPException(status_code=404, detail="Task not found or not active") + + mcp_server = FuzeAgentMCPServer() + mcp_integration = MCPClaudeIntegration(mcp_server) + + session_id = execution.claude_session_id or f"session-{task_id}" + context = await mcp_integration.get_session_context( + session_id=session_id, agent_id=execution.agent_id, task_id=task_id + ) + + return context + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get MCP context: {str(e)}" + ) + + +@app.post( + "/agents/{agent_id}/mcp-setup", + tags=["mcp-integration"], + summary="Setup Agent MCP Integration", + description="Configure MCP integration for an AI agent", +) +async def setup_agent_mcp( + agent_id: str = Path(..., description="Agent ID"), + setup_data: AgentMCPSetupRequest = Body(...), +): + """Set up MCP integration for an agent""" + try: + from .mcp_integration import FuzeAgentMCPServer, MCPClaudeIntegration + + task_id = setup_data.get("task_id") + session_id = setup_data.get("session_id") + + if not task_id: + raise HTTPException(status_code=400, detail="task_id is required") + + mcp_server = FuzeAgentMCPServer() + mcp_integration = MCPClaudeIntegration(mcp_server) + + # Set up MCP for Claude session + mcp_config = await mcp_integration.setup_claude_session_mcp( + session_id=session_id or f"session-{task_id}", + agent_id=agent_id, + task_id=task_id, + ) + + return { + "agent_id": agent_id, + "task_id": task_id, + "mcp_config": mcp_config, + "status": "mcp_configured", + } + + except Exception as e: + raise HTTPException(status_code=500, detail=f"Failed to setup MCP: {str(e)}") + + +# Multi-Agent Coordination Endpoints +@app.post( + "/tasks/{task_id}/coordinate", + tags=["multi-agent-coordination"], + summary="Initiate Multi-Agent Coordination", + description="Initiate multi-agent coordination for complex tasks", + response_model=CoordinationResponse, +) +async def initiate_task_coordination( + task_id: str = Path(..., description="Task ID to coordinate"), + coordination_request: CoordinationRequest = Body(...), +): + """Initiate multi-agent coordination for a complex task""" + try: + from .multi_agent_coordinator import CoordinationMode + + coordination_mode = coordination_request.get( + "coordination_mode", "collaborative" + ) + required_agents = coordination_request.get("required_agents") + required_skills = coordination_request.get("required_skills") + + # Validate coordination mode + try: + coord_mode = CoordinationMode(coordination_mode) + except ValueError: + raise HTTPException( + status_code=400, + detail=f"Invalid coordination mode: {coordination_mode}", + ) + + # Get multi-agent coordinator + coordinator = getattr( + app.state.task_execution_engine, "multi_agent_coordinator", None + ) + if not coordinator: + raise HTTPException( + status_code=503, detail="Multi-agent coordination not available" + ) + + # Initiate coordination + session_id = await coordinator.initiate_coordination( + task_id=task_id, + coordination_mode=coord_mode, + required_agents=required_agents, + required_skills=required_skills, + ) + + if session_id: + return { + "task_id": task_id, + "coordination_session_id": session_id, + "status": "coordination_initiated", + "coordination_mode": coordination_mode, + } + else: + return { + "task_id": task_id, + "status": "coordination_not_needed", + "message": "Task does not require multi-agent coordination", + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to initiate coordination: {str(e)}" + ) + + +@app.get("/coordination/{session_id}") +async def get_coordination_status(session_id: str): + """Get status of a coordination session""" + try: + coordinator = getattr( + app.state.task_execution_engine, "multi_agent_coordinator", None + ) + if not coordinator: + raise HTTPException( + status_code=503, detail="Multi-agent coordination not available" + ) + + status = await coordinator.get_coordination_status(session_id) + + if status: + return status + else: + raise HTTPException( + status_code=404, detail="Coordination session not found" + ) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get coordination status: {str(e)}" + ) + + +@app.post("/coordination/{session_id}/cancel") +async def cancel_coordination(session_id: str): + """Cancel a coordination session""" + try: + coordinator = getattr( + app.state.task_execution_engine, "multi_agent_coordinator", None + ) + if not coordinator: + raise HTTPException( + status_code=503, detail="Multi-agent coordination not available" + ) + + success = await coordinator.cancel_coordination(session_id) + + if success: + return {"coordination_session_id": session_id, "status": "cancelled"} + else: + raise HTTPException( + status_code=404, detail="Coordination session not found" + ) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to cancel coordination: {str(e)}" + ) + + +@app.post("/agents/{from_agent_id}/communicate/{to_agent_id}") +async def send_agent_communication( + from_agent_id: str, to_agent_id: str, communication_data: dict +): + """Send communication between agents""" + try: + message_type = communication_data.get("message_type", "notification") + content = communication_data.get("content", "") + metadata = communication_data.get("metadata", {}) + + if not content: + raise HTTPException(status_code=400, detail="Content cannot be empty") + + coordinator = getattr( + app.state.task_execution_engine, "multi_agent_coordinator", None + ) + if not coordinator: + raise HTTPException( + status_code=503, detail="Multi-agent coordination not available" + ) + + communication_id = await coordinator.send_agent_communication( + from_agent_id=from_agent_id, + to_agent_id=to_agent_id, + message_type=message_type, + content=content, + metadata=metadata, + ) + + return { + "communication_id": communication_id, + "from_agent_id": from_agent_id, + "to_agent_id": to_agent_id, + "status": "sent", + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to send agent communication: {str(e)}" + ) + + +@app.get("/coordination/active") +async def get_active_coordinations(): + """Get all active coordination sessions""" + try: + coordinator = getattr( + app.state.task_execution_engine, "multi_agent_coordinator", None + ) + if not coordinator: + raise HTTPException( + status_code=503, detail="Multi-agent coordination not available" + ) + + active_sessions = [] + for session_id in coordinator.active_sessions.keys(): + status = await coordinator.get_coordination_status(session_id) + if status: + active_sessions.append(status) + + return {"active_coordinations": active_sessions, "count": len(active_sessions)} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get active coordinations: {str(e)}" + ) + + +# WebSocket for coordination updates +@app.websocket("/ws/coordination/{session_id}") +async def coordination_websocket_endpoint(websocket: WebSocket, session_id: str): + """WebSocket endpoint for real-time coordination updates""" + await websocket.accept() + try: + coordinator = getattr( + app.state.task_execution_engine, "multi_agent_coordinator", None + ) + if not coordinator: + await websocket.send_json( + {"type": "error", "message": "Multi-agent coordination not available"} + ) + await websocket.close() + return + + # Monitor coordination session + while True: + try: + status = await coordinator.get_coordination_status(session_id) + if status: + await websocket.send_json( + { + "type": "coordination_update", + "session_id": session_id, + "data": status, + "timestamp": datetime.now().isoformat(), + } + ) + + # If coordination is completed or failed, send final update + if status.get("status") in ["completed", "failed", "cancelled"]: + await websocket.send_json( + { + "type": "coordination_finished", + "session_id": session_id, + "final_status": status.get("status"), + "timestamp": datetime.now().isoformat(), + } + ) + break + else: + await websocket.send_json( + { + "type": "error", + "message": f"Coordination session {session_id} not found", + } + ) + break + + await asyncio.sleep(3) # Update every 3 seconds + + except Exception as e: + await websocket.send_json( + { + "type": "error", + "message": f"Error monitoring coordination: {str(e)}", + } + ) + + except Exception as e: + print(f"Coordination WebSocket error for {session_id}: {e}") + finally: + await websocket.close() + + +# --------------------------------------------------------------------------- +# Agent relay WebSocket (Track 4) +# --------------------------------------------------------------------------- +@app.websocket("/agent-relay/{agent_id}") +async def agent_relay_endpoint(websocket: WebSocket, agent_id: str): + """ + Agent pods connect here to stream their session output. + Dashboard clients connect here to watch a specific agent's session. + Both use the same endpoint — first JSON message determines role: + {"role": "agent"} -> agent pod streaming output + {"role": "subscriber"} -> human dashboard watcher (default) + """ + await websocket.accept() + role = None + try: + init_msg = await websocket.receive_json() + role = init_msg.get("role", "subscriber") + + if role == "agent": + # Stream from agent pod to all subscribers + async for data in websocket.iter_json(): + msg = {"agentId": agent_id, **data} + dead = [] + for sub in list(agent_relay_subscribers[agent_id]): + try: + await sub.send_json(msg) + except Exception: + dead.append(sub) + for d in dead: + agent_relay_subscribers[agent_id].remove(d) + else: + # Human dashboard subscriber — wait for messages from agent + agent_relay_subscribers[agent_id].append(websocket) + await websocket.receive_text() # keep alive until disconnect + except WebSocketDisconnect: + pass + except Exception as e: + logger.warning(f"agent-relay {agent_id}: {e}") + finally: + subs = agent_relay_subscribers.get(agent_id, []) + if role != "agent" and websocket in subs: + subs.remove(websocket) + + +# Model Configuration and API Key Management Endpoints +@app.post( + "/organizations/{organization_id}/providers/{provider}/credentials", + tags=["model-configuration"], + summary="Store Provider API Credentials", + description="Store encrypted API credentials for a model provider", +) +async def store_provider_credentials( + organization_id: str = Path(..., description="Organization ID"), + provider: str = Path(..., description="Provider name"), + credentials: ProviderCredentialsRequest = Body(...), +): + """Store encrypted API credentials for a model provider at organization level""" + try: + from .model_configuration import ModelProvider, model_config_manager + + # Validate provider + try: + provider_enum = ModelProvider(provider) + except ValueError: + raise HTTPException( + status_code=400, detail=f"Unsupported provider: {provider}" + ) + + success = await model_config_manager.store_provider_credentials( + organization_id=organization_id, + provider=provider_enum, + api_key=credentials.api_key, + endpoint_url=credentials.endpoint_url, + additional_config=credentials.additional_config, + ) + + if success: + return { + "organization_id": organization_id, + "provider": provider, + "status": "credentials_stored", + "message": "API credentials stored successfully", + } + else: + raise HTTPException(status_code=500, detail="Failed to store credentials") + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to store provider credentials: {str(e)}" + ) + + +@app.get( + "/organizations/{organization_id}/models", + tags=["model-configuration"], + summary="Get Available Models", + description="Get available AI models for an organization", +) +async def get_available_models( + organization_id: str = Path(..., description="Organization ID"), + provider: Optional[str] = Query(None, description="Filter by provider"), + capabilities: Optional[str] = Query( + None, description="Filter by capabilities (comma-separated)" + ), +): + """Get available AI models with provider credential validation""" + try: + from .model_configuration import ( + ModelCapability, + ModelProvider, + model_config_manager, + ) + + provider_filter = None + if provider: + try: + provider_filter = ModelProvider(provider) + except ValueError: + raise HTTPException( + status_code=400, detail=f"Invalid provider: {provider}" + ) + + capabilities_filter = None + if capabilities: + try: + capabilities_filter = [ + ModelCapability(cap.strip()) for cap in capabilities.split(",") + ] + except ValueError as e: + raise HTTPException( + status_code=400, detail=f"Invalid capability: {str(e)}" + ) + + models = await model_config_manager.get_available_models( + organization_id=organization_id, + provider=provider_filter, + capabilities=capabilities_filter, + ) + + return { + "organization_id": organization_id, + "models": models, + "count": len(models), + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get available models: {str(e)}" + ) + + +@app.post( + "/agents/{agent_id}/model-configuration", + tags=["model-configuration"], + summary="Configure Agent Model Settings", + description="Configure model settings and preferences for an agent", +) +async def configure_agent_model( + agent_id: str = Path(..., description="Agent ID"), + config: AgentModelConfigRequest = Body(...), +): + """Configure model settings for an AI agent""" + try: + from .model_configuration import AgentModelConfig, model_config_manager + + agent_config = AgentModelConfig( + agent_id=agent_id, + primary_model=config.primary_model, + fallback_models=config.fallback_models, + temperature=config.temperature, + max_tokens=config.max_tokens, + top_p=config.top_p, + frequency_penalty=config.frequency_penalty, + presence_penalty=config.presence_penalty, + custom_instructions=config.custom_instructions, + use_function_calling=config.use_function_calling, + streaming_enabled=config.streaming_enabled, + cost_limit_per_task=config.cost_limit_per_task, + ) + + success = await model_config_manager.configure_agent_model( + agent_id, agent_config + ) + + if success: + return { + "agent_id": agent_id, + "status": "configured", + "primary_model": config.primary_model, + "fallback_models": config.fallback_models, + } + else: + raise HTTPException( + status_code=500, detail="Failed to configure agent model" + ) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to configure agent model: {str(e)}" + ) + + +@app.get( + "/agents/{agent_id}/model-configuration", + tags=["model-configuration"], + summary="Get Agent Model Configuration", + description="Get current model configuration for an agent", +) +async def get_agent_model_configuration( + agent_id: str = Path(..., description="Agent ID") +): + """Get model configuration for an AI agent""" + try: + from .model_configuration import model_config_manager + + config = await model_config_manager.get_agent_model_config(agent_id) + + if config: + return { + "agent_id": agent_id, + "configuration": { + "primary_model": config.primary_model, + "fallback_models": config.fallback_models, + "temperature": config.temperature, + "max_tokens": config.max_tokens, + "top_p": config.top_p, + "frequency_penalty": config.frequency_penalty, + "presence_penalty": config.presence_penalty, + "custom_instructions": config.custom_instructions, + "use_function_calling": config.use_function_calling, + "streaming_enabled": config.streaming_enabled, + "cost_limit_per_task": config.cost_limit_per_task, + "created_at": config.created_at.isoformat(), + "updated_at": config.updated_at.isoformat(), + }, + } + else: + raise HTTPException( + status_code=404, detail="Agent model configuration not found" + ) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get agent model configuration: {str(e)}" + ) + + +@app.post( + "/agents/{agent_id}/tasks/cost-estimate", + tags=["model-configuration"], + summary="Estimate Task Cost", + description="Estimate the cost of executing a task with the agent's model configuration", +) +async def estimate_task_cost( + agent_id: str = Path(..., description="Agent ID"), + request: TaskCostEstimateRequest = Body(...), +): + """Estimate cost for task execution based on agent's model configuration""" + try: + from .model_configuration import model_config_manager + + estimate = await model_config_manager.estimate_task_cost( + agent_id=agent_id, + task_description=request.task_description, + estimated_complexity=request.estimated_complexity, + ) + + return estimate + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to estimate task cost: {str(e)}" + ) + + +@app.get( + "/organizations/{organization_id}/model-usage", + tags=["model-configuration"], + summary="Get Model Usage Statistics", + description="Get model usage statistics and costs for an organization", +) +async def get_organization_model_usage( + organization_id: str = Path(..., description="Organization ID"), + days: int = Query(30, ge=1, le=365, description="Number of days to analyze"), +): + """Get model usage statistics and costs for an organization""" + try: + from .model_configuration import model_config_manager + + usage = await model_config_manager.get_organization_model_usage( + organization_id=organization_id, days=days + ) + + return usage + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get model usage: {str(e)}" + ) + + +@app.get( + "/agents/{agent_id}/model-recommendations", + tags=["model-configuration"], + summary="Get Model Recommendations", + description="Get model recommendations for an agent based on task capabilities", +) +async def get_model_recommendations( + agent_id: str = Path(..., description="Agent ID"), + capabilities: str = Query( + ..., description="Required capabilities (comma-separated)" + ), + cost_limit: Optional[float] = Query( + None, ge=0.0, description="Maximum cost limit in USD" + ), +): + """Get model recommendations based on task capabilities and cost constraints""" + try: + from .model_configuration import ModelCapability, model_config_manager + + # Parse capabilities + try: + capability_list = [ + ModelCapability(cap.strip()) for cap in capabilities.split(",") + ] + except ValueError as e: + raise HTTPException(status_code=400, detail=f"Invalid capability: {str(e)}") + + recommended_model = await model_config_manager.get_model_for_task( + agent_id=agent_id, task_capabilities=capability_list, cost_limit=cost_limit + ) + + if recommended_model: + return { + "agent_id": agent_id, + "recommended_model": recommended_model, + "capabilities": capabilities, + "cost_limit": cost_limit, + } + else: + return { + "agent_id": agent_id, + "recommended_model": None, + "message": "No suitable model found for the specified requirements", + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get model recommendations: {str(e)}" + ) + + +# Knowledge Management and Notification Endpoints + + +@app.get( + "/knowledge/notifications/{recipient_type}/{recipient_id}", + tags=["knowledge-management"], + summary="Get Knowledge Notifications", + description="Get notifications about knowledge updates, conflicts, and opportunities", +) +async def get_knowledge_notifications( + recipient_type: str = Path( + ..., description="Recipient type (agent, team, organization)" + ), + recipient_id: str = Path(..., description="Recipient ID"), + limit: int = Query(20, ge=1, le=100, description="Maximum notifications to return"), + status_filter: Optional[str] = Query( + None, description="Filter by status (unread, read, acknowledged)" + ), + notification_type_filter: Optional[str] = Query( + None, description="Filter by type (comma-separated)" + ), +): + """Get knowledge notifications for a recipient""" + try: + from .knowledge_notification_service import ( + KnowledgeNotificationService, + NotificationStatus, + NotificationType, + ) + + # Initialize notification service if not already done + if not hasattr(app.state, "notification_service"): + database_url = os.getenv( + "DATABASE_URL", + "postgresql://postgres:password@postgres:5432/ai_context", + ) + app.state.notification_service = KnowledgeNotificationService(database_url) + await app.state.notification_service.initialize() + + # Parse filters + status_filters = None + if status_filter: + try: + status_filters = [ + NotificationStatus(s.strip()) for s in status_filter.split(",") + ] + except ValueError as e: + raise HTTPException( + status_code=400, detail=f"Invalid status filter: {str(e)}" + ) + + type_filters = None + if notification_type_filter: + try: + type_filters = [ + NotificationType(t.strip()) + for t in notification_type_filter.split(",") + ] + except ValueError as e: + raise HTTPException( + status_code=400, + detail=f"Invalid notification type filter: {str(e)}", + ) + + notifications = ( + await app.state.notification_service.get_notifications_for_recipient( + recipient_type=recipient_type, + recipient_id=recipient_id, + limit=limit, + status_filter=status_filters, + notification_type_filter=type_filters, + ) + ) + + return { + "recipient_type": recipient_type, + "recipient_id": recipient_id, + "notifications": [ + { + "id": n.id, + "notification_type": n.notification_type.value, + "title": n.title, + "message": n.message, + "knowledge_id": n.knowledge_id, + "knowledge_type": n.knowledge_type, + "priority": n.priority.value, + "requires_action": n.requires_action, + "status": n.status.value, + "suggested_actions": n.suggested_actions, + "metadata": n.metadata, + "created_at": n.created_at.isoformat(), + "expires_at": n.expires_at.isoformat() if n.expires_at else None, + } + for n in notifications + ], + "count": len(notifications), + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get knowledge notifications: {str(e)}" + ) + + +@app.put( + "/knowledge/notifications/{notification_id}/status", + tags=["knowledge-management"], + summary="Update Notification Status", + description="Mark notification as read, acknowledged, or acted upon", +) +async def update_notification_status( + notification_id: str = Path(..., description="Notification ID"), + status: str = Body(..., description="New notification status"), + action_taken: Optional[Dict[str, Any]] = Body( + None, description="Optional action taken metadata" + ), +): + """Update notification status and optional action taken""" + try: + from .knowledge_notification_service import NotificationStatus + + # Validate status + try: + notification_status = NotificationStatus(status) + except ValueError: + raise HTTPException(status_code=400, detail=f"Invalid status: {status}") + + success = await app.state.notification_service.mark_notification_status( + notification_id=notification_id, + status=notification_status, + action_taken=action_taken, + ) + + if success: + return { + "notification_id": notification_id, + "status": status, + "updated": True, + } + else: + raise HTTPException(status_code=404, detail="Notification not found") + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to update notification status: {str(e)}" + ) + + +@app.get( + "/knowledge/notifications/statistics", + tags=["knowledge-management"], + summary="Get Notification Statistics", + description="Get comprehensive notification statistics and analytics", +) +async def get_notification_statistics( + organization_id: Optional[str] = Query( + None, description="Filter by organization ID" + ), + days_back: int = Query(30, ge=1, le=365, description="Days of history to analyze"), +): + """Get notification statistics and analytics""" + try: + stats = await app.state.notification_service.get_notification_statistics( + organization_id=organization_id, days_back=days_back + ) + + return stats + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get notification statistics: {str(e)}" + ) + + +@app.post( + "/knowledge/organizations/{organization_id}/add", + tags=["knowledge-management"], + summary="Add Organizational Knowledge", + description="Add knowledge to organization-level knowledge base", +) +async def add_organizational_knowledge( + organization_id: str = Path(..., description="Organization ID"), + title: str = Body(..., description="Knowledge title"), + content: str = Body(..., description="Knowledge content"), + content_type: str = Body("documentation", description="Content type"), + knowledge_category: str = Body("development", description="Knowledge category"), + source_agent_id: Optional[str] = Body(None, description="Source agent ID"), + source_team_id: Optional[str] = Body(None, description="Source team ID"), + tags: List[str] = Body(default_factory=list, description="Knowledge tags"), + metadata: Dict[str, Any] = Body( + default_factory=dict, description="Additional metadata" + ), +): + """Add knowledge to organization-level knowledge base""" + try: + from .organization_rag_manager import ( + ContentType, + KnowledgeCategory, + OrganizationRAGManager, + SourceType, + ) + + # Initialize services if not already done + if not hasattr(app.state, "org_rag_manager"): + database_url = os.getenv( + "DATABASE_URL", + "postgresql://postgres:password@postgres:5432/ai_context", + ) + app.state.org_rag_manager = OrganizationRAGManager(database_url) + await app.state.org_rag_manager.initialize() + + # Validate enums + try: + content_type_enum = ContentType(content_type) + category_enum = KnowledgeCategory(knowledge_category) + except ValueError as e: + raise HTTPException(status_code=400, detail=f"Invalid enum value: {str(e)}") + + knowledge_id = await app.state.org_rag_manager.add_knowledge( + organization_id=organization_id, + title=title, + content=content, + content_type=content_type_enum, + knowledge_category=category_enum, + source_type=SourceType.MANUAL_INPUT, + source_agent_id=source_agent_id, + source_team_id=source_team_id, + tags=tags, + metadata=metadata, + ) + + return { + "knowledge_id": knowledge_id, + "organization_id": organization_id, + "title": title, + "status": "added", + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to add organizational knowledge: {str(e)}" + ) + + +@app.get( + "/knowledge/organizations/{organization_id}/search", + tags=["knowledge-management"], + summary="Search Organizational Knowledge", + description="Search organization-level knowledge base", +) +async def search_organizational_knowledge( + organization_id: str = Path(..., description="Organization ID"), + query: str = Query(..., description="Search query"), + limit: int = Query(10, ge=1, le=50, description="Maximum results to return"), + min_similarity: float = Query( + 0.3, ge=0.0, le=1.0, description="Minimum similarity threshold" + ), + categories: Optional[str] = Query( + None, description="Filter by categories (comma-separated)" + ), +): + """Search organization-level knowledge base""" + try: + from .organization_rag_manager import KnowledgeCategory + + # Parse categories + category_filters = None + if categories: + try: + category_filters = [ + KnowledgeCategory(cat.strip()) for cat in categories.split(",") + ] + except ValueError as e: + raise HTTPException( + status_code=400, detail=f"Invalid category: {str(e)}" + ) + + search_results = await app.state.org_rag_manager.search_knowledge( + organization_id=organization_id, + query=query, + limit=limit, + min_similarity=min_similarity, + categories=category_filters, + ) + + results = [] + for result in search_results: + results.append( + { + "knowledge_id": result.knowledge.id, + "title": result.knowledge.title, + "content_preview": ( + result.knowledge.content[:200] + "..." + if len(result.knowledge.content) > 200 + else result.knowledge.content + ), + "category": result.knowledge.knowledge_category.value, + "content_type": result.knowledge.content_type.value, + "similarity_score": result.similarity_score, + "combined_score": result.combined_score, + "quality_score": result.knowledge.quality_score, + "usage_count": result.knowledge.usage_count, + "created_at": result.knowledge.created_at.isoformat(), + "tags": result.knowledge.tags, + "metadata": result.knowledge.metadata, + } + ) + + return { + "organization_id": organization_id, + "query": query, + "results": results, + "count": len(results), + } + + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"Failed to search organizational knowledge: {str(e)}", + ) + + +@app.get( + "/knowledge/context-enhancement/{agent_id}", + tags=["knowledge-management"], + summary="Get Enhanced Context for Agent", + description="Get enhanced context with relevant organizational knowledge for task execution", +) +async def get_enhanced_context_for_agent( + agent_id: str = Path(..., description="Agent ID"), + task_description: str = Query( + ..., description="Task description for context enhancement" + ), + task_type: Optional[str] = Query(None, description="Task type"), + technologies: Optional[str] = Query( + None, description="Technologies involved (comma-separated)" + ), +): + """Get enhanced context with relevant knowledge for agent task execution""" + try: + from .context_enhancement_service import ContextEnhancementService + + # Initialize context enhancement service if needed + if not hasattr(app.state, "context_enhancement_service"): + database_url = os.getenv( + "DATABASE_URL", + "postgresql://postgres:password@postgres:5432/ai_context", + ) + # These would be initialized in the lifespan + if hasattr(app.state, "org_rag_manager") and hasattr( + app.state, "team_knowledge_manager" + ): + app.state.context_enhancement_service = ContextEnhancementService( + database_url=database_url, + org_rag_manager=app.state.org_rag_manager, + team_knowledge_manager=app.state.team_knowledge_manager, + ) + await app.state.context_enhancement_service.initialize() + else: + raise HTTPException( + status_code=503, + detail="Knowledge management services not initialized", + ) + + # Build task data + task_data = { + "description": task_description, + "task_type": task_type, + "technologies": technologies.split(",") if technologies else [], + } + + enhanced_context = ( + await app.state.context_enhancement_service.enhance_agent_context( + agent_id=agent_id, task_data=task_data + ) + ) + + return { + "agent_id": agent_id, + "task_description": task_description, + "enhanced_context": { + "organizational_knowledge_count": len( + enhanced_context.organizational_knowledge + ), + "team_knowledge_count": len(enhanced_context.team_knowledge), + "similar_task_insights_count": len( + enhanced_context.similar_task_insights + ), + "success_patterns": enhanced_context.success_patterns, + "common_pitfalls": enhanced_context.common_pitfalls, + "recommended_approaches": enhanced_context.recommended_approaches, + "context_summary": enhanced_context.context_summary, + "enhancement_metadata": enhanced_context.enhancement_metadata, + }, + "organizational_knowledge": [ + { + "knowledge_id": item.knowledge_id, + "title": item.title, + "category": item.category, + "relevance_score": item.relevance_score, + "confidence_score": item.confidence_score, + "content_preview": ( + item.content[:200] + "..." + if len(item.content) > 200 + else item.content + ), + } + for item in enhanced_context.organizational_knowledge + ], + "team_knowledge": [ + { + "knowledge_id": item.knowledge_id, + "title": item.title, + "category": item.category, + "relevance_score": item.relevance_score, + "confidence_score": item.confidence_score, + "content_preview": ( + item.content[:200] + "..." + if len(item.content) > 200 + else item.content + ), + } + for item in enhanced_context.team_knowledge + ], + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get enhanced context: {str(e)}" + ) + + +@app.get( + "/knowledge/analytics/organizations/{organization_id}/insights", + tags=["knowledge-management"], + summary="Get Organizational Knowledge Insights", + description="Get comprehensive analytics and insights about organizational knowledge", +) +async def get_organizational_knowledge_insights( + organization_id: str = Path(..., description="Organization ID"), + analysis_period_days: int = Query( + 30, ge=7, le=365, description="Analysis period in days" + ), +): + """Get comprehensive organizational knowledge insights and analytics""" + try: + insights = ( + await app.state.knowledge_analytics_service.get_organizational_insights( + organization_id=organization_id, + analysis_period_days=analysis_period_days, + ) + ) + + return { + "organization_id": organization_id, + "analysis_period_days": analysis_period_days, + "insights": { + "total_knowledge_items": insights.total_knowledge_items, + "knowledge_growth_rate": insights.knowledge_growth_rate, + "knowledge_utilization_rate": insights.knowledge_utilization_rate, + "knowledge_freshness_score": insights.knowledge_freshness_score, + "cross_team_sharing_rate": insights.cross_team_sharing_rate, + "propagation_efficiency": insights.propagation_efficiency, + "top_performing_categories": insights.top_performing_categories, + "knowledge_gaps": insights.knowledge_gaps, + "agent_knowledge_engagement": insights.agent_knowledge_engagement, + "team_knowledge_contribution": insights.team_knowledge_contribution, + "recommendations": insights.recommendations, + }, + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get organizational insights: {str(e)}" + ) + + +@app.get( + "/knowledge/analytics/organizations/{organization_id}/effectiveness", + tags=["knowledge-management"], + summary="Analyze Knowledge Effectiveness", + description="Analyze effectiveness and performance of knowledge items", +) +async def analyze_knowledge_effectiveness( + organization_id: str = Path(..., description="Organization ID"), + knowledge_category: Optional[str] = Query( + None, description="Filter by knowledge category" + ), + min_usage_count: int = Query( + 3, ge=1, description="Minimum usage count for analysis" + ), +): + """Analyze effectiveness of knowledge items in the organization""" + try: + effectiveness_metrics = ( + await app.state.knowledge_analytics_service.analyze_knowledge_effectiveness( + organization_id=organization_id, + knowledge_category=knowledge_category, + min_usage_count=min_usage_count, + ) + ) + + results = [] + for metric in effectiveness_metrics: + results.append( + { + "knowledge_id": metric.knowledge_id, + "title": metric.title, + "category": metric.category, + "usage_count": metric.usage_count, + "success_correlation": metric.success_correlation, + "average_relevance": metric.average_relevance, + "agent_adoption_rate": metric.agent_adoption_rate, + "team_adoption_rate": metric.team_adoption_rate, + "quality_score": metric.quality_score, + "recency_score": metric.recency_score, + "overall_effectiveness": metric.overall_effectiveness, + "trend_direction": metric.trend_direction, + "optimization_suggestions": metric.optimization_suggestions, + } + ) + + return { + "organization_id": organization_id, + "effectiveness_analysis": results, + "total_analyzed": len(results), + "summary": { + "avg_effectiveness": sum(r["overall_effectiveness"] for r in results) + / max(len(results), 1), + "top_performers": sorted( + results, key=lambda x: x["overall_effectiveness"], reverse=True + )[:5], + "needs_attention": [ + r for r in results if r["overall_effectiveness"] < 0.5 + ], + }, + } + + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"Failed to analyze knowledge effectiveness: {str(e)}", + ) + + +@app.get( + "/knowledge/analytics/agents/{agent_id}/profile", + tags=["knowledge-management"], + summary="Get Agent Knowledge Profile", + description="Get detailed knowledge profile and analytics for an agent", +) +async def get_agent_knowledge_profile( + agent_id: str = Path(..., description="Agent ID"), + analysis_period_days: int = Query( + 60, ge=7, le=365, description="Analysis period in days" + ), +): + """Get detailed knowledge profile for an agent""" + try: + profile = ( + await app.state.knowledge_analytics_service.get_agent_knowledge_profile( + agent_id=agent_id, analysis_period_days=analysis_period_days + ) + ) + + if not profile: + raise HTTPException( + status_code=404, detail="Agent not found or no knowledge data available" + ) + + return { + "agent_id": agent_id, + "analysis_period_days": analysis_period_days, + "profile": { + "agent_name": profile.agent_name, + "team_id": profile.team_id, + "knowledge_consumption_rate": profile.knowledge_consumption_rate, + "knowledge_creation_rate": profile.knowledge_creation_rate, + "expertise_areas": profile.expertise_areas, + "knowledge_application_success": profile.knowledge_application_success, + "learning_velocity": profile.learning_velocity, + "knowledge_sharing_activity": profile.knowledge_sharing_activity, + "preferred_knowledge_types": profile.preferred_knowledge_types, + "knowledge_gaps": profile.knowledge_gaps, + "optimization_recommendations": profile.optimization_recommendations, + }, + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get agent knowledge profile: {str(e)}" + ) + + +@app.get( + "/knowledge/analytics/organizations/{organization_id}/optimization", + tags=["knowledge-management"], + summary="Get Knowledge Optimization Recommendations", + description="Get comprehensive recommendations for knowledge system optimization", +) +async def get_knowledge_optimization_recommendations( + organization_id: str = Path(..., description="Organization ID"), + focus_area: Optional[str] = Query( + None, + description="Focus area (utilization, quality, gaps, propagation, collaboration)", + ), +): + """Generate comprehensive knowledge optimization recommendations""" + try: + recommendations = await app.state.knowledge_analytics_service.generate_knowledge_optimization_recommendations( + organization_id=organization_id, focus_area=focus_area + ) + + return { + "organization_id": organization_id, + "focus_area": focus_area, + "recommendations": recommendations, + "total_recommendations": len(recommendations), + } + + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"Failed to get optimization recommendations: {str(e)}", + ) + + +@app.get( + "/knowledge/analytics/organizations/{organization_id}/trends", + tags=["knowledge-management"], + summary="Get Knowledge Trends Analysis", + description="Analyze knowledge trends and patterns over time", +) +async def get_knowledge_trends_analysis( + organization_id: str = Path(..., description="Organization ID"), + trend_period_days: int = Query( + 90, ge=30, le=365, description="Trend analysis period in days" + ), +): + """Get comprehensive knowledge trends analysis""" + try: + trends = ( + await app.state.knowledge_analytics_service.get_knowledge_trends_analysis( + organization_id=organization_id, trend_period_days=trend_period_days + ) + ) + + return { + "organization_id": organization_id, + "trend_period_days": trend_period_days, + "trends": trends, + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get knowledge trends: {str(e)}" + ) + + +# Memory-Enhanced Agents Endpoints + + +@app.post( + "/agents/{agent_id}/deploy-memory", + tags=["memory-agents"], + summary="Deploy Memory-Enabled Agent", + description="Deploy an agent with persistent memory capabilities", +) +async def deploy_memory_enabled_agent( + agent_id: str = Path(..., description="Agent ID"), + template_id: str = Body(..., description="Agent template ID"), + task_id: Optional[str] = Body(None, description="Optional specific task ID"), + repository_settings: Optional[Dict[str, Any]] = Body( + None, description="Repository settings" + ), +): + """Deploy a memory-enabled autonomous agent container""" + try: + result = await app.state.agent_manager.deploy_memory_enabled_agent( + agent_id=agent_id, + template_id=template_id, + task_id=task_id, + repository_settings=repository_settings, + ) + + if result["success"]: + return result + else: + raise HTTPException(status_code=500, detail=result["error"]) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to deploy memory-enabled agent: {str(e)}" + ) + + +@app.get( + "/agents/{agent_id}/memory-status", + tags=["memory-agents"], + summary="Get Agent Memory Status", + description="Get agent memory status and expertise summary", +) +async def get_agent_memory_status(agent_id: str = Path(..., description="Agent ID")): + """Get agent memory status, expertise metrics, and insights""" + try: + status = await app.state.agent_manager.get_agent_memory_status(agent_id) + return status + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get agent memory status: {str(e)}" + ) + + +@app.post( + "/agents/{agent_id}/memory-tasks", + tags=["memory-agents"], + summary="Assign Task to Memory Agent", + description="Assign a task to a memory-enabled agent", +) +async def assign_task_to_memory_agent( + agent_id: str = Path(..., description="Agent ID"), + task_id: str = Body(..., description="Task ID"), + task_data: Dict[str, Any] = Body(..., description="Task data"), +): + """Assign a task to a memory-enabled agent for autonomous execution""" + try: + result = await app.state.agent_manager.assign_task_to_memory_agent( + agent_id=agent_id, task_id=task_id, task_data=task_data + ) + + if result["success"]: + return result + else: + raise HTTPException(status_code=400, detail=result["error"]) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to assign task to memory agent: {str(e)}" + ) + + +@app.delete( + "/agents/{agent_id}/memory", + tags=["memory-agents"], + summary="Stop Memory-Enabled Agent", + description="Stop a memory-enabled agent container", +) +async def stop_memory_enabled_agent(agent_id: str = Path(..., description="Agent ID")): + """Stop and clean up a memory-enabled agent container""" + try: + result = await app.state.agent_manager.stop_memory_enabled_agent(agent_id) + + if result["success"]: + return result + else: + raise HTTPException(status_code=400, detail=result["error"]) + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to stop memory-enabled agent: {str(e)}" + ) + + +@app.get( + "/system/expertise-dashboard", + tags=["memory-agents"], + summary="Get System Expertise Dashboard", + description="Get system-wide expertise and memory analytics", +) +async def get_system_expertise_dashboard(): + """Get comprehensive dashboard of system expertise and memory analytics""" + try: + dashboard = await app.state.agent_manager.get_system_expertise_dashboard() + return dashboard + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get expertise dashboard: {str(e)}" + ) + + +@app.get( + "/agents/{agent_id}/tasks/pending", + tags=["memory-agents"], + summary="Get Pending Tasks for Agent", + description="Get pending tasks for a memory-enabled agent", +) +async def get_pending_tasks_for_agent( + agent_id: str = Path(..., description="Agent ID"), + limit: int = Query( + 10, ge=1, le=50, description="Maximum number of tasks to return" + ), +): + """Get pending tasks that a memory-enabled agent can pick up""" + try: + async with get_db_connection() as conn: + tasks = await conn.fetch( + """ + SELECT id, title, description, type, complexity, language, + framework, requirements, created_at + FROM tasks + WHERE agent_id = $1 + AND status = 'pending' + AND assigned_to_memory_agent = true + ORDER BY created_at ASC + LIMIT $2 + """, + agent_id, + limit, + ) + + return { + "agent_id": agent_id, + "tasks": [dict(task) for task in tasks], + "count": len(tasks), + } + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to get pending tasks: {str(e)}" + ) + + +@app.put( + "/tasks/{task_id}/status", + tags=["memory-agents"], + summary="Update Task Status", + description="Update task status (used by memory-enabled agents)", +) +async def update_task_status( + task_id: str = Path(..., description="Task ID"), + status: str = Body(..., description="New task status"), + result: Optional[Dict[str, Any]] = Body(None, description="Task result data"), + updated_by: Optional[str] = Body(None, description="ID of agent updating the task"), + container_instance_id: Optional[str] = Body( + None, description="Container instance ID" + ), + updated_at: Optional[str] = Body(None, description="Update timestamp"), +): + """Update task status - used by memory-enabled agents to report progress""" + try: + async with get_db_connection() as conn: + await conn.execute( + """ + UPDATE tasks + SET status = $2, + result = COALESCE($3, result), + updated_by = COALESCE($4, updated_by), + updated_at = NOW() + WHERE id = $1 + """, + task_id, + status, + result, + updated_by, + ) + + # If task is completed, log it for expertise tracking + if status in ["completed", "failed"]: + # The agent's memory system will handle learning from the outcome + pass + + return {"task_id": task_id, "status": status, "updated": True} + + except Exception as e: + raise HTTPException( + status_code=500, detail=f"Failed to update task status: {str(e)}" + ) + + +@app.post( + "/agents/{agent_id}/register", + tags=["memory-agents"], + summary="Agent Registration", + description="Register agent capabilities and status with orchestrator", +) +async def register_agent_capabilities( + agent_id: str = Path(..., description="Agent ID"), + capabilities: Dict[str, Any] = Body( + ..., description="Agent capabilities and status" + ), +): + """Register or update agent capabilities - used by memory-enabled agents on startup""" + try: + # Update agent capabilities in database + async with get_db_connection() as conn: + await conn.execute( + """ + UPDATE agents + SET config = config || $2, + status = 'active', + updated_at = NOW() + WHERE id = $1 + """, + agent_id, + { + "capabilities": capabilities, + "last_registration": datetime.now().isoformat(), + }, + ) + + # Update in-memory tracking + if agent_id in app.state.agent_manager.memory_enabled_agents: + app.state.agent_manager.memory_enabled_agents[agent_id]["status"] = "active" + + return { + "agent_id": agent_id, + "agent_recognized": True, + "capabilities_accepted": True, + "status": "registered", + } + + except Exception as e: + return { + "agent_id": agent_id, + "agent_recognized": False, + "capabilities_accepted": False, + "error": str(e), + } + + +@app.post( + "/agents/{agent_id}/statistics", + tags=["memory-agents"], + summary="Agent Statistics Update", + description="Update agent performance and memory statistics", +) +async def update_agent_statistics( + agent_id: str = Path(..., description="Agent ID"), + stats: Dict[str, Any] = Body(..., description="Agent statistics"), +): + """Update agent statistics - used by memory-enabled agents for performance tracking""" + try: + # Store statistics for analytics + async with get_db_connection() as conn: + await conn.execute( + """ + UPDATE agents + SET config = config || $2, + updated_at = NOW() + WHERE id = $1 + """, + agent_id, + { + "latest_statistics": stats, + "statistics_updated_at": datetime.now().isoformat(), + }, + ) + + # Clear expertise cache to force refresh + await app.state.agent_manager.expertise_tracker.clear_cache(agent_id) + + return {"agent_id": agent_id, "statistics_updated": True} + + except Exception as e: + return {"agent_id": agent_id, "statistics_updated": False, "error": str(e)} + + +@app.post( + "/agents/{agent_id}/error", + tags=["memory-agents"], + summary="Agent Error Reporting", + description="Report agent errors for monitoring", +) +async def report_agent_error( + agent_id: str = Path(..., description="Agent ID"), + error_data: Dict[str, Any] = Body(..., description="Error information"), +): + """Report agent errors - used by memory-enabled agents for error tracking""" + try: + # Log error for monitoring + logger.error(f"Agent {agent_id} reported error: {error_data}") + + # Update agent status if it's a critical error + if error_data.get("critical", False): + async with get_db_connection() as conn: + await conn.execute( + """ + UPDATE agents + SET status = 'error', + config = config || $2, + updated_at = NOW() + WHERE id = $1 + """, + agent_id, + { + "last_error": error_data, + "error_reported_at": datetime.now().isoformat(), + }, + ) + + return {"agent_id": agent_id, "error_logged": True} + + except Exception as e: + logger.error(f"Failed to log agent error: {e}") + return {"agent_id": agent_id, "error_logged": False} + + +# ============================================================================ +# Goals Management API Endpoints +# ============================================================================ + + +@app.post( + "/organizations/{organization_id}/goals", + tags=["goals-management"], + summary="Create organizational goal", + description="Create a new goal for an organization with specified targets and deadlines", +) +async def create_goal( + organization_id: str = Path(..., description="Organization ID"), + goal_data: GoalCreateRequest = Body(..., description="Goal creation data"), + created_by: Optional[str] = Query( + None, description="ID of user/agent creating the goal" + ), +): + """Create a new organizational goal""" + try: + from .goals_management_service import GoalType + + goal_id = await app.state.goals_service.create_goal( + organization_id=organization_id, + title=goal_data.title, + description=goal_data.description, + goal_type=GoalType(goal_data.goal_type), + target_value=goal_data.target_value, + target_unit=goal_data.target_unit, + target_deadline=goal_data.target_deadline, + priority_level=goal_data.priority_level, + success_criteria=goal_data.success_criteria, + assigned_teams=goal_data.assigned_teams, + goal_owner_agent_id=goal_data.goal_owner_agent_id, + stakeholder_agents=goal_data.stakeholder_agents, + tags=goal_data.tags, + metadata=goal_data.metadata, + created_by=created_by, + ) + + return {"goal_id": goal_id, "status": "created"} + + except Exception as e: + logger.error(f"Error creating goal: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/organizations/{organization_id}/goals", + tags=["goals-management"], + summary="List organization goals", + description="Get all goals for an organization with optional filtering", +) +async def list_organization_goals( + organization_id: str = Path(..., description="Organization ID"), + status: Optional[List[str]] = Query(None, description="Filter by goal status"), + goal_type: Optional[List[str]] = Query(None, description="Filter by goal type"), + limit: int = Query( + 50, ge=1, le=100, description="Maximum number of goals to return" + ), +): + """List goals for an organization""" + try: + from .goals_management_service import GoalStatus, GoalType + + status_filter = [GoalStatus(s) for s in status] if status else None + type_filter = [GoalType(gt) for gt in goal_type] if goal_type else None + + goals = await app.state.goals_service.list_organization_goals( + organization_id=organization_id, + status_filter=status_filter, + goal_type_filter=type_filter, + limit=limit, + ) + + return { + "organization_id": organization_id, + "goals": [ + { + "id": goal.id, + "title": goal.title, + "description": goal.description, + "goal_type": goal.goal_type.value, + "status": goal.status.value, + "progress_percentage": float(goal.progress_percentage), + "target_value": ( + float(goal.target_value) if goal.target_value else None + ), + "target_unit": goal.target_unit, + "current_value": ( + float(goal.current_value) if goal.current_value else None + ), + "target_deadline": goal.target_deadline.isoformat(), + "priority_level": goal.priority_level, + "completion_confidence": float(goal.completion_confidence), + "created_at": goal.created_at.isoformat(), + "updated_at": goal.updated_at.isoformat(), + } + for goal in goals + ], + } + + except Exception as e: + logger.error(f"Error listing organization goals: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/goals/{goal_id}", + tags=["goals-management"], + summary="Get goal details", + description="Get detailed information about a specific goal", +) +async def get_goal(goal_id: str = Path(..., description="Goal ID")): + """Get goal details""" + try: + goal = await app.state.goals_service.get_goal(goal_id) + + if not goal: + raise HTTPException(status_code=404, detail="Goal not found") + + return { + "id": goal.id, + "organization_id": goal.organization_id, + "title": goal.title, + "description": goal.description, + "goal_type": goal.goal_type.value, + "status": goal.status.value, + "progress_percentage": float(goal.progress_percentage), + "target_value": float(goal.target_value) if goal.target_value else None, + "target_unit": goal.target_unit, + "current_value": float(goal.current_value) if goal.current_value else None, + "success_criteria": goal.success_criteria, + "start_date": goal.start_date.isoformat(), + "target_deadline": goal.target_deadline.isoformat(), + "actual_completion_date": ( + goal.actual_completion_date.isoformat() + if goal.actual_completion_date + else None + ), + "priority_level": goal.priority_level, + "completion_confidence": float(goal.completion_confidence), + "assigned_teams": goal.assigned_teams, + "goal_owner_agent_id": goal.goal_owner_agent_id, + "stakeholder_agents": goal.stakeholder_agents, + "tags": goal.tags, + "metadata": goal.metadata, + "created_by": goal.created_by, + "created_at": goal.created_at.isoformat(), + "updated_at": goal.updated_at.isoformat(), + } + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting goal {goal_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/goals/{goal_id}/overview", + tags=["goals-management"], + summary="Get goal overview", + description="Get comprehensive overview of goal with milestones, tasks, and progress", +) +async def get_goal_overview(goal_id: str = Path(..., description="Goal ID")): + """Get comprehensive goal overview""" + try: + overview = await app.state.goals_service.get_goal_overview(goal_id) + + if not overview: + raise HTTPException(status_code=404, detail="Goal not found") + + return overview + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting goal overview {goal_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.put( + "/goals/{goal_id}/progress", + tags=["goals-management"], + summary="Update goal progress", + description="Update progress for a specific goal", +) +async def update_goal_progress( + goal_id: str = Path(..., description="Goal ID"), + progress_data: GoalUpdateRequest = Body(..., description="Progress update data"), + recorded_by: Optional[str] = Query( + None, description="ID of user/agent recording progress" + ), +): + """Update goal progress""" + try: + success = await app.state.goals_service.update_goal_progress( + goal_id=goal_id, + progress_percentage=progress_data.progress_percentage, + current_value=progress_data.current_value, + completion_confidence=progress_data.completion_confidence, + progress_notes=progress_data.notes, + recorded_by=recorded_by, + ) + + if not success: + raise HTTPException( + status_code=404, detail="Goal not found or no changes made" + ) + + return {"goal_id": goal_id, "status": "updated"} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error updating goal progress {goal_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/goals/{goal_id}/milestones", + tags=["goals-management"], + summary="Create milestone", + description="Create a new milestone for a goal", +) +async def create_milestone( + goal_id: str = Path(..., description="Goal ID"), + milestone_data: MilestoneCreateRequest = Body( + ..., description="Milestone creation data" + ), + created_by: Optional[str] = Query( + None, description="ID of user/agent creating milestone" + ), +): + """Create milestone for goal""" + try: + milestone_id = await app.state.goals_service.create_milestone( + goal_id=goal_id, + title=milestone_data.title, + description=milestone_data.description, + target_date=milestone_data.target_date, + milestone_type=milestone_data.milestone_type, + success_criteria=milestone_data.success_criteria, + deliverables=milestone_data.deliverables, + dependencies=milestone_data.dependencies, + assigned_teams=milestone_data.assigned_teams, + responsible_agent_id=milestone_data.responsible_agent_id, + priority_level=milestone_data.priority_level, + weight_in_goal=milestone_data.weight_in_goal, + created_by=created_by, + ) + + return {"milestone_id": milestone_id, "status": "created"} + + except Exception as e: + logger.error(f"Error creating milestone: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/milestones/{milestone_id}/tasks", + tags=["goals-management"], + summary="Create task from milestone", + description="Create a new task derived from a milestone", +) +async def create_task_from_milestone( + milestone_id: str = Path(..., description="Milestone ID"), + task_data: TaskFromMilestoneRequest = Body(..., description="Task creation data"), + created_by: Optional[str] = Query( + None, description="ID of user/agent creating task" + ), +): + """Create task from milestone""" + try: + task_id = await app.state.goals_service.create_task_from_milestone( + milestone_id=milestone_id, + title=task_data.title, + description=task_data.description, + task_type=task_data.task_type, + complexity_level=task_data.complexity_level, + estimated_hours=task_data.estimated_hours, + due_date=task_data.due_date, + assigned_team_id=task_data.assigned_team_id, + assigned_agent_id=task_data.assigned_agent_id, + priority=task_data.priority, + requirements=task_data.requirements, + acceptance_criteria=task_data.acceptance_criteria, + dependencies=task_data.dependencies, + created_by_agent_id=created_by, + ) + + return {"task_id": task_id, "status": "created"} + + except Exception as e: + logger.error(f"Error creating task from milestone: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/goals/{goal_id}/generate-execution-plan", + tags=["goals-management"], + summary="Generate execution plan", + description="Generate comprehensive milestone and task execution plan for a goal", +) +async def generate_execution_plan( + goal_id: str = Path(..., description="Goal ID"), + planning_context: Optional[Dict[str, Any]] = Body( + None, description="Additional planning context" + ), +): + """Generate execution plan with milestones and tasks""" + try: + execution_plan = ( + await app.state.milestone_task_engine.generate_goal_execution_plan( + goal_id=goal_id, planning_context=planning_context + ) + ) + + return execution_plan + + except Exception as e: + logger.error(f"Error generating execution plan for goal {goal_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/goals/{goal_id}/generate-monthly-milestones", + tags=["goals-management"], + summary="Generate monthly milestones", + description="Generate monthly milestone breakdown for a goal", +) +async def generate_monthly_milestones( + goal_id: str = Path(..., description="Goal ID"), + start_date: Optional[date] = Query(None, description="Start date for milestones"), + end_date: Optional[date] = Query(None, description="End date for milestones"), +): + """Generate monthly milestones for goal""" + try: + milestone_ids = ( + await app.state.milestone_task_engine.generate_monthly_milestones( + goal_id=goal_id, start_date=start_date, end_date=end_date + ) + ) + + return { + "goal_id": goal_id, + "milestone_ids": milestone_ids, + "count": len(milestone_ids), + "status": "generated", + } + + except Exception as e: + logger.error(f"Error generating monthly milestones for goal {goal_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/milestones/{milestone_id}/generate-weekly-tasks", + tags=["goals-management"], + summary="Generate weekly tasks", + description="Generate weekly task breakdown for a milestone", +) +async def generate_weekly_tasks( + milestone_id: str = Path(..., description="Milestone ID"), + focus_areas: Optional[List[str]] = Body( + None, description="Focus areas for task generation" + ), +): + """Generate weekly tasks for milestone""" + try: + task_ids = ( + await app.state.milestone_task_engine.generate_weekly_tasks_for_milestone( + milestone_id=milestone_id, focus_areas=focus_areas + ) + ) + + return { + "milestone_id": milestone_id, + "task_ids": task_ids, + "count": len(task_ids), + "status": "generated", + } + + except Exception as e: + logger.error(f"Error generating weekly tasks for milestone {milestone_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/goals/{goal_id}/generate-cross-functional-tasks", + tags=["goals-management"], + summary="Generate cross-functional tasks", + description="Generate tasks across different business functions for a goal", +) +async def generate_cross_functional_tasks( + goal_id: str = Path(..., description="Goal ID"), + target_functions: Optional[List[str]] = Body( + None, description="Target business functions" + ), +): + """Generate cross-functional tasks for goal""" + try: + functional_tasks = ( + await app.state.milestone_task_engine.generate_cross_functional_tasks( + goal_id=goal_id, target_functions=target_functions + ) + ) + + return { + "goal_id": goal_id, + "functional_tasks": functional_tasks, + "total_tasks": sum(len(tasks) for tasks in functional_tasks.values()), + "status": "generated", + } + + except Exception as e: + logger.error(f"Error generating cross-functional tasks for goal {goal_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/goals/{goal_id}/conversations", + tags=["goals-management"], + summary="Create goal conversation", + description="Create AI-powered conversation for goal planning and discussion", +) +async def create_goal_conversation( + goal_id: str = Path(..., description="Goal ID"), + conversation_data: GoalConversationCreateRequest = Body( + ..., description="Conversation creation data" + ), + created_by: Optional[str] = Query( + None, description="ID of user/agent creating conversation" + ), +): + """Create goal conversation""" + try: + from .goal_conversation_service import ConversationType + + conversation_id = ( + await app.state.goal_conversation_service.create_goal_conversation( + goal_id=goal_id, + conversation_type=ConversationType(conversation_data.conversation_type), + conversation_title=conversation_data.conversation_title, + initial_context=conversation_data.initial_context, + participants=conversation_data.participants, + created_by=created_by, + ) + ) + + return {"conversation_id": conversation_id, "status": "created"} + + except Exception as e: + logger.error(f"Error creating goal conversation: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/conversations/{conversation_id}", + tags=["goals-management"], + summary="Get goal conversation", + description="Get full conversation with messages, insights, and action items", +) +async def get_goal_conversation( + conversation_id: str = Path(..., description="Conversation ID") +): + """Get goal conversation""" + try: + conversation = await app.state.goal_conversation_service.get_conversation( + conversation_id + ) + + if not conversation: + raise HTTPException(status_code=404, detail="Conversation not found") + + return conversation + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting conversation {conversation_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/conversations/{conversation_id}/messages", + tags=["goals-management"], + summary="Add message to conversation", + description="Add a new message to a goal conversation", +) +async def add_message_to_conversation( + conversation_id: str = Path(..., description="Conversation ID"), + message_data: ConversationMessageRequest = Body(..., description="Message data"), + sender_id: Optional[str] = Query(None, description="ID of message sender"), +): + """Add message to conversation""" + try: + from .goal_conversation_service import MessageType + + message_id = ( + await app.state.goal_conversation_service.add_message_to_conversation( + conversation_id=conversation_id, + message_type=MessageType(message_data.message_type), + sender_id=sender_id, + sender_name=message_data.sender_name, + content=message_data.content, + metadata=message_data.metadata, + references=message_data.references, + ) + ) + + return {"message_id": message_id, "status": "added"} + + except Exception as e: + logger.error(f"Error adding message to conversation: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/conversations/{conversation_id}/generate-milestones", + tags=["goals-management"], + summary="Generate milestones from conversation", + description="Generate milestone recommendations based on conversation analysis", +) +async def generate_planning_milestones( + conversation_id: str = Path(..., description="Conversation ID"), + planning_context: Optional[Dict[str, Any]] = Body( + None, description="Additional planning context" + ), +): + """Generate planning milestones from conversation""" + try: + milestones = ( + await app.state.goal_conversation_service.generate_planning_milestones( + conversation_id=conversation_id, planning_context=planning_context + ) + ) + + return { + "conversation_id": conversation_id, + "milestones": milestones, + "count": len(milestones), + "status": "generated", + } + + except Exception as e: + logger.error(f"Error generating planning milestones: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/conversations/{conversation_id}/conduct-progress-review", + tags=["goals-management"], + summary="Conduct progress review", + description="Conduct AI-powered progress review for a goal conversation", +) +async def conduct_progress_review( + conversation_id: str = Path(..., description="Conversation ID"), + review_period_days: int = Query( + 30, ge=1, le=365, description="Review period in days" + ), +): + """Conduct progress review""" + try: + review_analysis = ( + await app.state.goal_conversation_service.conduct_progress_review( + conversation_id=conversation_id, review_period_days=review_period_days + ) + ) + + return review_analysis + + except Exception as e: + logger.error(f"Error conducting progress review: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/conversations/{conversation_id}/extract-action-items", + tags=["goals-management"], + summary="Extract action items", + description="Extract and create action items from conversation analysis", +) +async def extract_action_items( + conversation_id: str = Path(..., description="Conversation ID"), + auto_assign: bool = Query( + True, description="Whether to automatically assign action items" + ), +): + """Extract action items from conversation""" + try: + action_items = await app.state.goal_conversation_service.extract_action_items_from_conversation( + conversation_id=conversation_id, auto_assign=auto_assign + ) + + return { + "conversation_id": conversation_id, + "action_items": action_items, + "count": len(action_items), + "status": "extracted", + } + + except Exception as e: + logger.error(f"Error extracting action items: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/goals/{goal_id}/conversations", + tags=["goals-management"], + summary="Get goal conversations", + description="Get all conversations for a goal with optional filtering", +) +async def get_goal_conversations( + goal_id: str = Path(..., description="Goal ID"), + conversation_type: Optional[str] = Query( + None, description="Filter by conversation type" + ), + status: Optional[str] = Query(None, description="Filter by conversation status"), + limit: int = Query(10, ge=1, le=50, description="Maximum number of conversations"), +): + """Get conversations for a goal""" + try: + from .goal_conversation_service import ConversationStatus, ConversationType + + conv_type = ConversationType(conversation_type) if conversation_type else None + conv_status = ConversationStatus(status) if status else None + + conversations = ( + await app.state.goal_conversation_service.get_goal_conversations( + goal_id=goal_id, + conversation_type=conv_type, + status=conv_status, + limit=limit, + ) + ) + + return { + "goal_id": goal_id, + "conversations": conversations, + "count": len(conversations), + } + + except Exception as e: + logger.error(f"Error getting goal conversations: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/goals/{goal_id}/track-progress", + tags=["goals-management"], + summary="Record progress tracking update", + description="Record detailed progress update with tracking and risk assessment", +) +async def record_progress_tracking( + goal_id: str = Path(..., description="Goal ID"), + progress_data: ProgressUpdateRequest = Body( + ..., description="Progress tracking data" + ), + recorded_by: Optional[str] = Query( + None, description="ID of user/agent recording progress" + ), +): + """Record progress tracking update""" + try: + snapshot_id = await app.state.goal_tracking_service.record_progress_update( + goal_id=goal_id, + progress_percentage=progress_data.progress_percentage, + current_value=progress_data.current_value, + milestone_id=progress_data.milestone_id, + notes=progress_data.notes, + recorded_by=recorded_by, + confidence_score=progress_data.confidence_score, + trigger_alerts=progress_data.trigger_alerts, + ) + + return {"goal_id": goal_id, "snapshot_id": snapshot_id, "status": "recorded"} + + except Exception as e: + logger.error(f"Error recording progress tracking: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/goals/{goal_id}/deadline-risk", + tags=["goals-management"], + summary="Assess deadline risk", + description="Get comprehensive deadline risk assessment for a goal", +) +async def assess_deadline_risk(goal_id: str = Path(..., description="Goal ID")): + """Assess deadline risk for goal""" + try: + deadline_risk = await app.state.goal_tracking_service.assess_goal_deadline_risk( + goal_id + ) + + return { + "goal_id": deadline_risk.goal_id, + "risk_level": deadline_risk.risk_level.value, + "probability_of_delay": float(deadline_risk.probability_of_delay), + "estimated_completion_date": deadline_risk.estimated_completion_date.isoformat(), + "days_at_risk": deadline_risk.days_at_risk, + "critical_path_items": deadline_risk.critical_path_items, + "mitigation_strategies": deadline_risk.mitigation_strategies, + "updated_at": deadline_risk.updated_at.isoformat(), + } + + except Exception as e: + logger.error(f"Error assessing deadline risk: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/goals/{goal_id}/progress-report", + tags=["goals-management"], + summary="Generate progress report", + description="Generate comprehensive progress report for a goal", +) +async def generate_progress_report( + goal_id: str = Path(..., description="Goal ID"), + report_period_days: int = Query( + 30, ge=1, le=365, description="Report period in days" + ), +): + """Generate progress report for goal""" + try: + report = await app.state.goal_tracking_service.generate_progress_report( + goal_id=goal_id, report_period_days=report_period_days + ) + + return report + + except Exception as e: + logger.error(f"Error generating progress report: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/organizations/{organization_id}/goals-dashboard", + tags=["goals-management"], + summary="Get organization goals dashboard", + description="Get comprehensive dashboard for all organization goals", +) +async def get_organization_goals_dashboard( + organization_id: str = Path(..., description="Organization ID") +): + """Get organization goals dashboard""" + try: + dashboard = await app.state.goals_service.get_organization_goals_dashboard( + organization_id + ) + return dashboard + + except Exception as e: + logger.error(f"Error getting organization dashboard: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/organizations/{organization_id}/tracking-dashboard", + tags=["goals-management"], + summary="Get tracking dashboard", + description="Get comprehensive tracking dashboard with risk assessments", +) +async def get_tracking_dashboard( + organization_id: str = Path(..., description="Organization ID") +): + """Get organization tracking dashboard""" + try: + dashboard = ( + await app.state.goal_tracking_service.get_organization_tracking_dashboard( + organization_id + ) + ) + return dashboard + + except Exception as e: + logger.error(f"Error getting tracking dashboard: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# ================================ +# Knowledge Management API Endpoints +# ================================ + + +@app.post( + "/knowledge/organizations/{organization_id}/documents", + tags=["knowledge-management"], + summary="Upload Organizational Document", + response_model=DocumentMetadata, +) +async def upload_organization_document( + organization_id: str = Path(..., description="Organization ID"), + file: UploadFile = File(..., description="Document file to upload"), + title: Optional[str] = Form(None, description="Document title"), + tags: Optional[str] = Form(None, description="Comma-separated tags"), +): + """Upload a document to organizational knowledge base""" + try: + tags_list = [] + if tags: + tags_list = [tag.strip() for tag in tags.split(",")] + + document = await knowledge_manager.upload_document( + file_content=file.file, + filename=file.filename, + title=title, + organization_id=organization_id, + tags=tags_list, + ) + + return document + + except Exception as e: + logger.error(f"Error uploading organizational document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/knowledge/organizations/{organization_id}/url", + tags=["knowledge-management"], + summary="Add URL to Organizational Knowledge", + response_model=DocumentMetadata, +) +async def add_organization_url( + organization_id: str = Path(..., description="Organization ID"), + url: str = Body(..., embed=True), + title: Optional[str] = Body(None, embed=True), + tags: Optional[List[str]] = Body(None, embed=True), +): + """Add URL content to organizational knowledge base""" + try: + document = await knowledge_manager.upload_url( + url=url, title=title, organization_id=organization_id, tags=tags or [] + ) + + return document + + except Exception as e: + logger.error(f"Error adding organizational URL: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/organizations/{organization_id}/documents", + tags=["knowledge-management"], + summary="List Organizational Documents", + response_model=List[DocumentMetadata], +) +async def list_organization_documents( + organization_id: str = Path(..., description="Organization ID") +): + """Get list of organizational documents""" + try: + documents = await knowledge_manager.get_documents( + organization_id=organization_id + ) + return documents + + except Exception as e: + logger.error(f"Error listing organizational documents: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/organizations/{organization_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Get Organizational Document", + response_model=DocumentMetadata, +) +async def get_organization_document( + organization_id: str = Path(..., description="Organization ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Get organizational document metadata""" + try: + document = await knowledge_manager.get_document_metadata( + doc_id=doc_id, organization_id=organization_id + ) + + if not document: + raise HTTPException(status_code=404, detail="Document not found") + + return document + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting organizational document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/organizations/{organization_id}/documents/{doc_id}/content", + tags=["knowledge-management"], + summary="Get Organizational Document Content", +) +async def get_organization_document_content( + organization_id: str = Path(..., description="Organization ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Get full content of organizational document""" + try: + content = await knowledge_manager.get_document_content( + doc_id=doc_id, organization_id=organization_id + ) + + if content is None: + raise HTTPException(status_code=404, detail="Document not found") + + return {"content": content} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting organizational document content: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.put( + "/knowledge/organizations/{organization_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Update Organizational Document", + response_model=DocumentMetadata, +) +async def update_organization_document( + organization_id: str = Path(..., description="Organization ID"), + doc_id: str = Path(..., description="Document ID"), + title: Optional[str] = Body(None, embed=True), + tags: Optional[List[str]] = Body(None, embed=True), +): + """Update organizational document metadata""" + try: + document = await knowledge_manager.update_document( + doc_id=doc_id, title=title, tags=tags, organization_id=organization_id + ) + + if not document: + raise HTTPException(status_code=404, detail="Document not found") + + return document + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error updating organizational document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.delete( + "/knowledge/organizations/{organization_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Delete Organizational Document", +) +async def delete_organization_document( + organization_id: str = Path(..., description="Organization ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Delete organizational document""" + try: + success = await knowledge_manager.delete_document( + doc_id=doc_id, organization_id=organization_id + ) + + if not success: + raise HTTPException(status_code=404, detail="Document not found") + + return {"message": "Document deleted successfully"} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error deleting organizational document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# Team Knowledge Management Endpoints + + +@app.post( + "/knowledge/teams/{team_id}/documents", + tags=["knowledge-management"], + summary="Upload Team Document", + response_model=DocumentMetadata, +) +async def upload_team_document( + team_id: str = Path(..., description="Team ID"), + file: UploadFile = File(..., description="Document file to upload"), + title: Optional[str] = Form(None, description="Document title"), + tags: Optional[str] = Form(None, description="Comma-separated tags"), +): + """Upload a document to team knowledge base""" + try: + tags_list = [] + if tags: + tags_list = [tag.strip() for tag in tags.split(",")] + + document = await knowledge_manager.upload_document( + file_content=file.file, + filename=file.filename, + title=title, + team_id=team_id, + tags=tags_list, + ) + + return document + + except Exception as e: + logger.error(f"Error uploading team document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/knowledge/teams/{team_id}/url", + tags=["knowledge-management"], + summary="Add URL to Team Knowledge", + response_model=DocumentMetadata, +) +async def add_team_url( + team_id: str = Path(..., description="Team ID"), + url: str = Body(..., embed=True), + title: Optional[str] = Body(None, embed=True), + tags: Optional[List[str]] = Body(None, embed=True), +): + """Add URL content to team knowledge base""" + try: + document = await knowledge_manager.upload_url( + url=url, title=title, team_id=team_id, tags=tags or [] + ) + + return document + + except Exception as e: + logger.error(f"Error adding team URL: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/teams/{team_id}/documents", + tags=["knowledge-management"], + summary="List Team Documents", + response_model=List[DocumentMetadata], +) +async def list_team_documents(team_id: str = Path(..., description="Team ID")): + """Get list of team documents""" + try: + documents = await knowledge_manager.get_documents(team_id=team_id) + return documents + + except Exception as e: + logger.error(f"Error listing team documents: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/teams/{team_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Get Team Document", + response_model=DocumentMetadata, +) +async def get_team_document( + team_id: str = Path(..., description="Team ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Get team document metadata""" + try: + document = await knowledge_manager.get_document_metadata( + doc_id=doc_id, team_id=team_id + ) + + if not document: + raise HTTPException(status_code=404, detail="Document not found") + + return document + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting team document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/teams/{team_id}/documents/{doc_id}/content", + tags=["knowledge-management"], + summary="Get Team Document Content", +) +async def get_team_document_content( + team_id: str = Path(..., description="Team ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Get full content of team document""" + try: + content = await knowledge_manager.get_document_content( + doc_id=doc_id, team_id=team_id + ) + + if content is None: + raise HTTPException(status_code=404, detail="Document not found") + + return {"content": content} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting team document content: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.put( + "/knowledge/teams/{team_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Update Team Document", + response_model=DocumentMetadata, +) +async def update_team_document( + team_id: str = Path(..., description="Team ID"), + doc_id: str = Path(..., description="Document ID"), + title: Optional[str] = Body(None, embed=True), + tags: Optional[List[str]] = Body(None, embed=True), +): + """Update team document metadata""" + try: + document = await knowledge_manager.update_document( + doc_id=doc_id, title=title, tags=tags, team_id=team_id + ) + + if not document: + raise HTTPException(status_code=404, detail="Document not found") + + return document + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error updating team document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.delete( + "/knowledge/teams/{team_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Delete Team Document", +) +async def delete_team_document( + team_id: str = Path(..., description="Team ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Delete team document""" + try: + success = await knowledge_manager.delete_document( + doc_id=doc_id, team_id=team_id + ) + + if not success: + raise HTTPException(status_code=404, detail="Document not found") + + return {"message": "Document deleted successfully"} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error deleting team document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# Agent Knowledge Management Endpoints + + +@app.post( + "/knowledge/agents/{agent_id}/documents", + tags=["knowledge-management"], + summary="Upload Agent Document", + response_model=DocumentMetadata, +) +async def upload_agent_document( + agent_id: str = Path(..., description="Agent ID"), + file: UploadFile = File(..., description="Document file to upload"), + title: Optional[str] = Form(None, description="Document title"), + tags: Optional[str] = Form(None, description="Comma-separated tags"), +): + """Upload a document to agent knowledge base""" + try: + tags_list = [] + if tags: + tags_list = [tag.strip() for tag in tags.split(",")] + + document = await knowledge_manager.upload_document( + file_content=file.file, + filename=file.filename, + title=title, + agent_id=agent_id, + tags=tags_list, + ) + + return document + + except Exception as e: + logger.error(f"Error uploading agent document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/knowledge/agents/{agent_id}/url", + tags=["knowledge-management"], + summary="Add URL to Agent Knowledge", + response_model=DocumentMetadata, +) +async def add_agent_url( + agent_id: str = Path(..., description="Agent ID"), + url: str = Body(..., embed=True), + title: Optional[str] = Body(None, embed=True), + tags: Optional[List[str]] = Body(None, embed=True), +): + """Add URL content to agent knowledge base""" + try: + document = await knowledge_manager.upload_url( + url=url, title=title, agent_id=agent_id, tags=tags or [] + ) + + return document + + except Exception as e: + logger.error(f"Error adding agent URL: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/agents/{agent_id}/documents", + tags=["knowledge-management"], + summary="List Agent Documents", + response_model=List[DocumentMetadata], +) +async def list_agent_documents(agent_id: str = Path(..., description="Agent ID")): + """Get list of agent documents""" + try: + documents = await knowledge_manager.get_documents(agent_id=agent_id) + return documents + + except Exception as e: + logger.error(f"Error listing agent documents: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/agents/{agent_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Get Agent Document", + response_model=DocumentMetadata, +) +async def get_agent_document( + agent_id: str = Path(..., description="Agent ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Get agent document metadata""" + try: + document = await knowledge_manager.get_document_metadata( + doc_id=doc_id, agent_id=agent_id + ) + + if not document: + raise HTTPException(status_code=404, detail="Document not found") + + return document + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting agent document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/knowledge/agents/{agent_id}/documents/{doc_id}/content", + tags=["knowledge-management"], + summary="Get Agent Document Content", +) +async def get_agent_document_content( + agent_id: str = Path(..., description="Agent ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Get full content of agent document""" + try: + content = await knowledge_manager.get_document_content( + doc_id=doc_id, agent_id=agent_id + ) + + if content is None: + raise HTTPException(status_code=404, detail="Document not found") + + return {"content": content} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error getting agent document content: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.put( + "/knowledge/agents/{agent_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Update Agent Document", + response_model=DocumentMetadata, +) +async def update_agent_document( + agent_id: str = Path(..., description="Agent ID"), + doc_id: str = Path(..., description="Document ID"), + title: Optional[str] = Body(None, embed=True), + tags: Optional[List[str]] = Body(None, embed=True), +): + """Update agent document metadata""" + try: + document = await knowledge_manager.update_document( + doc_id=doc_id, title=title, tags=tags, agent_id=agent_id + ) + + if not document: + raise HTTPException(status_code=404, detail="Document not found") + + return document + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error updating agent document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.delete( + "/knowledge/agents/{agent_id}/documents/{doc_id}", + tags=["knowledge-management"], + summary="Delete Agent Document", +) +async def delete_agent_document( + agent_id: str = Path(..., description="Agent ID"), + doc_id: str = Path(..., description="Document ID"), +): + """Delete agent document""" + try: + success = await knowledge_manager.delete_document( + doc_id=doc_id, agent_id=agent_id + ) + + if not success: + raise HTTPException(status_code=404, detail="Document not found") + + return {"message": "Document deleted successfully"} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error deleting agent document: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# Knowledge Search Endpoints + + +@app.get( + "/knowledge/search", + tags=["knowledge-management"], + summary="Search Knowledge Base", + response_model=List[DocumentMetadata], +) +async def search_knowledge( + query: str = Query(..., description="Search query"), + organization_id: Optional[str] = Query(None, description="Filter by organization"), + team_id: Optional[str] = Query(None, description="Filter by team"), + agent_id: Optional[str] = Query(None, description="Filter by agent"), + limit: int = Query(10, ge=1, le=100, description="Maximum number of results"), +): + """Search across knowledge base""" + try: + documents = await knowledge_manager.search_documents( + query=query, + organization_id=organization_id, + team_id=team_id, + agent_id=agent_id, + limit=limit, + ) + + return documents + + except Exception as e: + logger.error(f"Error searching knowledge: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# ================================ +# Container Management API Endpoints +# ================================ + + +@app.post( + "/agents/{agent_id}/container/create", + tags=["container-management"], + summary="Create Agent Container", + response_model=ContainerStatus, +) +async def create_agent_container( + agent_id: str = Path(..., description="Agent ID"), + config: Optional[ContainerConfig] = Body( + None, description="Container configuration" + ), +): + """Create a new container for an AI agent""" + try: + status = await container_manager.create_agent_container(agent_id, config) + return status + + except Exception as e: + logger.error(f"Error creating container for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/agents/{agent_id}/container/start", + tags=["container-management"], + summary="Start Agent Container", + response_model=ContainerStatus, +) +async def start_agent_container(agent_id: str = Path(..., description="Agent ID")): + """Start an agent container""" + try: + status = await container_manager.start_container(agent_id) + return status + + except RuntimeError as e: + raise HTTPException(status_code=404, detail=str(e)) + except Exception as e: + logger.error(f"Error starting container for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/agents/{agent_id}/container/stop", + tags=["container-management"], + summary="Stop Agent Container", + response_model=ContainerStatus, +) +async def stop_agent_container( + agent_id: str = Path(..., description="Agent ID"), + timeout: int = Body(30, description="Stop timeout in seconds"), +): + """Stop an agent container""" + try: + status = await container_manager.stop_container(agent_id, timeout) + return status + + except RuntimeError as e: + raise HTTPException(status_code=404, detail=str(e)) + except Exception as e: + logger.error(f"Error stopping container for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/agents/{agent_id}/container/restart", + tags=["container-management"], + summary="Restart Agent Container", + response_model=ContainerStatus, +) +async def restart_agent_container( + agent_id: str = Path(..., description="Agent ID"), + timeout: int = Body(30, description="Restart timeout in seconds"), +): + """Restart an agent container""" + try: + status = await container_manager.restart_container(agent_id, timeout) + return status + + except RuntimeError as e: + raise HTTPException(status_code=404, detail=str(e)) + except Exception as e: + logger.error(f"Error restarting container for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.delete( + "/agents/{agent_id}/container", + tags=["container-management"], + summary="Remove Agent Container", +) +async def remove_agent_container( + agent_id: str = Path(..., description="Agent ID"), + force: bool = Query(False, description="Force removal of running container"), +): + """Remove an agent container""" + try: + success = await container_manager.remove_container(agent_id, force) + + if success: + return {"message": f"Container for agent {agent_id} removed successfully"} + else: + raise HTTPException(status_code=500, detail="Failed to remove container") + + except Exception as e: + logger.error(f"Error removing container for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/agents/{agent_id}/container/status", + tags=["container-management"], + summary="Get Agent Container Status", + response_model=Optional[ContainerStatus], +) +async def get_agent_container_status(agent_id: str = Path(..., description="Agent ID")): + """Get container status for an agent""" + try: + status = await container_manager.get_container_status(agent_id) + return status + + except Exception as e: + logger.error(f"Error getting container status for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/containers/agents", + tags=["container-management"], + summary="List Agent Containers", + response_model=List[ContainerStatus], +) +async def list_agent_containers(): + """List all agent containers""" + try: + containers = await container_manager.list_agent_containers() + return containers + + except Exception as e: + logger.error(f"Error listing agent containers: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get( + "/agents/{agent_id}/container/logs", + tags=["container-management"], + summary="Get Agent Container Logs", +) +async def get_agent_container_logs( + agent_id: str = Path(..., description="Agent ID"), + tail: int = Query(100, ge=1, le=10000, description="Number of log lines to return"), + since: Optional[str] = Query( + None, description="Show logs since timestamp (ISO format)" + ), +): + """Get container logs for an agent""" + try: + since_dt = None + if since: + try: + since_dt = datetime.fromisoformat(since.replace("Z", "+00:00")) + except ValueError: + raise HTTPException(status_code=400, detail="Invalid timestamp format") + + logs = await container_manager.get_container_logs( + agent_id=agent_id, tail=tail, since=since_dt + ) + + return {"logs": logs} + + except RuntimeError as e: + raise HTTPException(status_code=404, detail=str(e)) + except Exception as e: + logger.error(f"Error getting container logs for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post( + "/agents/{agent_id}/container/execute", + tags=["container-management"], + summary="Execute Command in Container", +) +async def execute_container_command( + agent_id: str = Path(..., description="Agent ID"), + command: str = Body(..., description="Command to execute"), + working_dir: Optional[str] = Body(None, description="Working directory"), +): + """Execute a command in the agent container""" + try: + result = await container_manager.execute_command( + agent_id=agent_id, command=command, working_dir=working_dir + ) + + return result + + except RuntimeError as e: + raise HTTPException(status_code=404, detail=str(e)) + except Exception as e: + logger.error(f"Error executing command in container for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# WebSocket endpoint for real-time log streaming +@app.websocket("/agents/{agent_id}/container/logs/stream") +async def stream_agent_container_logs(websocket: WebSocket, agent_id: str): + """Stream container logs in real-time via WebSocket""" + await websocket.accept() + + try: + # Check if container exists + status = await container_manager.get_container_status(agent_id) + if not status: + await websocket.send_json({"error": "Container not found"}) + await websocket.close() + return + + await websocket.send_json({"status": "connected", "agent_id": agent_id}) + + # Stream logs + async for log_entry in container_manager.stream_container_logs(agent_id): + await websocket.send_json( + { + "timestamp": log_entry.timestamp.isoformat(), + "stream": log_entry.stream, + "message": log_entry.message, + } + ) + + except Exception as e: + logger.error(f"Error in log stream for agent {agent_id}: {e}") + try: + await websocket.send_json({"error": str(e)}) + except Exception: + logger.debug("Failed to send error frame on closing websocket") + finally: + try: + await websocket.close() + except Exception: + logger.debug("Failed to close websocket cleanly") + + +# ============================================================================ +# RAG (Retrieval-Augmented Generation) Endpoints +# ============================================================================ + + +@app.post("/rag/search") +async def search_knowledge_context( + query: str = Body(..., embed=True), + organization_id: Optional[str] = Body(None, embed=True), + team_id: Optional[str] = Body(None, embed=True), + agent_id: Optional[str] = Body(None, embed=True), + max_results: int = Body(5, embed=True), + similarity_threshold: float = Body(0.7, embed=True), +): + """Search for relevant knowledge context using RAG""" + try: + context = await rag_system.search_relevant_context( + query=query, + organization_id=organization_id, + team_id=team_id, + agent_id=agent_id, + max_results=max_results, + similarity_threshold=similarity_threshold, + ) + + return { + "query": context.query, + "relevant_chunks": [ + { + "document_id": chunk.document_id, + "document_title": chunk.metadata.get("document_title", "Unknown"), + "content": chunk.content, + "chunk_index": chunk.chunk_index, + "metadata": chunk.metadata, + } + for chunk in context.relevant_chunks + ], + "similarity_scores": context.similarity_scores, + "total_documents": context.total_documents, + "context_length": context.context_length, + } + + except Exception as e: + logger.error(f"Error searching knowledge context: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post("/rag/enhance-prompt") +async def enhance_prompt_with_context( + message: str = Body(..., embed=True), + organization_id: Optional[str] = Body(None, embed=True), + team_id: Optional[str] = Body(None, embed=True), + agent_id: Optional[str] = Body(None, embed=True), + max_context_length: int = Body(4000, embed=True), +): + """Enhance a prompt with relevant context using RAG""" + try: + enhanced_prompt = await rag_system.get_contextual_prompt( + user_message=message, + organization_id=organization_id, + team_id=team_id, + agent_id=agent_id, + max_context_length=max_context_length, + ) + + return { + "original_message": message, + "enhanced_prompt": enhanced_prompt, + "context_added": len(enhanced_prompt) > len(message), + } + + except Exception as e: + logger.error(f"Error enhancing prompt with context: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post("/rag/reindex") +async def reindex_knowledge_base( + organization_id: Optional[str] = Body(None, embed=True), + team_id: Optional[str] = Body(None, embed=True), + agent_id: Optional[str] = Body(None, embed=True), +): + """Reindex all documents in a scope for RAG""" + try: + results = await rag_system.index_all_documents( + organization_id=organization_id, team_id=team_id, agent_id=agent_id + ) + + return { + "scope": { + "organization_id": organization_id, + "team_id": team_id, + "agent_id": agent_id, + }, + "results": results, + "message": f"Indexed {results['indexed']} documents, {results['failed']} failed, {results['skipped']} skipped", + } + + except Exception as e: + logger.error(f"Error reindexing knowledge base: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.get("/rag/stats") +async def get_rag_index_stats(): + """Get statistics about the RAG index""" + try: + stats = await rag_system.get_index_stats() + return stats + + except Exception as e: + logger.error(f"Error getting RAG stats: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@app.post("/rag/documents/{doc_id}/reindex") +async def reindex_document( + doc_id: str, + organization_id: Optional[str] = Body(None, embed=True), + team_id: Optional[str] = Body(None, embed=True), + agent_id: Optional[str] = Body(None, embed=True), +): + """Reindex a specific document for RAG""" + try: + # Get document metadata + document = await knowledge_manager.get_document_metadata( + doc_id=doc_id, + organization_id=organization_id, + team_id=team_id, + agent_id=agent_id, + ) + + if not document: + raise HTTPException(status_code=404, detail="Document not found") + + # Reindex the document + success = await rag_system.index_document(document) + + if success: + return { + "document_id": doc_id, + "status": "reindexed", + "message": f"Document '{document.title}' has been reindexed successfully", + } + else: + raise HTTPException(status_code=500, detail="Failed to reindex document") + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error reindexing document {doc_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# ============================================================================ +# Real-time WebSocket Endpoints +# ============================================================================ + + +@app.websocket("/ws/updates") +async def websocket_real_time_updates( + websocket: WebSocket, + organization_id: Optional[str] = None, + team_id: Optional[str] = None, + agent_id: Optional[str] = None, + user_id: Optional[str] = None, + subscriptions: Optional[str] = None, +): + """Main WebSocket endpoint for real-time updates""" + import uuid + + connection_id = str(uuid.uuid4()) + + # Parse subscriptions + subscription_list = [] + if subscriptions: + subscription_list = subscriptions.split(",") + + try: + connection = await websocket_manager.connect( + websocket=websocket, + connection_id=connection_id, + organization_id=organization_id, + team_id=team_id, + agent_id=agent_id, + user_id=user_id, + subscriptions=subscription_list, + ) + + # Keep connection alive and handle pings + while True: + try: + # Wait for ping messages or disconnection + message = await websocket.receive_text() + + # Handle ping/pong + if message == "ping": + await websocket.send_text("pong") + connection.last_ping = datetime.now() + else: + # Parse other messages (subscription updates, etc.) + try: + data = json.loads(message) + if data.get("type") == "subscribe": + # Update subscriptions + new_subs = data.get("subscriptions", []) + connection.scope.subscriptions.clear() + for sub in new_subs: + try: + connection.scope.subscriptions.add(UpdateType(sub)) + except ValueError: + pass + + await connection.send_update( + WebSocketUpdate( + type=UpdateType.SYSTEM_NOTIFICATION, + data={ + "message": "Subscriptions updated", + "subscriptions": list( + connection.scope.subscriptions + ), + }, + ) + ) + except json.JSONDecodeError: + pass + + except WebSocketDisconnect: + break + + except Exception as e: + logger.error(f"WebSocket error for connection {connection_id}: {e}") + finally: + await websocket_manager.disconnect(connection_id) + + +@app.websocket("/ws/agent/{agent_id}/updates") +async def websocket_agent_updates(websocket: WebSocket, agent_id: str): + """WebSocket endpoint for specific agent updates""" + import uuid + + connection_id = f"agent-{agent_id}-{uuid.uuid4()}" + + try: + connection = await websocket_manager.connect( + websocket=websocket, + connection_id=connection_id, + agent_id=agent_id, + subscriptions=[ + UpdateType.AGENT_STATUS.value, + UpdateType.TASK_STATUS.value, + UpdateType.TASK_PROGRESS.value, + UpdateType.CONTAINER_STATUS.value, + UpdateType.CHAT_MESSAGE.value, + UpdateType.CHAT_TYPING.value, + ], + ) + + # Keep connection alive + while True: + try: + message = await websocket.receive_text() + if message == "ping": + await websocket.send_text("pong") + connection.last_ping = datetime.now() + except WebSocketDisconnect: + break + + except Exception as e: + logger.error(f"Agent WebSocket error for {agent_id}: {e}") + finally: + await websocket_manager.disconnect(connection_id) + + +@app.websocket("/ws/agents/{agent_id}/conversations/{conversation_id}") +async def websocket_agent_conversation( + websocket: WebSocket, agent_id: str, conversation_id: str +): + """WebSocket endpoint for real-time agent conversation""" + import uuid + + connection_id = f"conversation-{conversation_id}-{uuid.uuid4()}" + + await websocket.accept() + + try: + # Store connection for broadcasting + active_conversations = getattr(app.state, "active_conversations", {}) + if conversation_id not in active_conversations: + active_conversations[conversation_id] = [] + active_conversations[conversation_id].append(websocket) + app.state.active_conversations = active_conversations + + while True: + try: + # Receive message from client + data = await websocket.receive_json() + + if data.get("type") == "ping": + await websocket.send_json({"type": "pong"}) + elif data.get("type") == "message": + # Handle new message + message_content = data.get("content", "") + if message_content: + # Store message in database + async with get_db_connection() as conn: + message_id = await conn.fetchval( + """ + INSERT INTO agent_conversations (session_id, agent_id, message_type, content) + VALUES ($1, $2, 'user', $3) + RETURNING id + """, + conversation_id, + agent_id, + message_content, + ) + + # Update session activity + await conn.execute( + """ + UPDATE chat_sessions + SET last_activity = CURRENT_TIMESTAMP, message_count = message_count + 1 + WHERE id = $1 + """, + conversation_id, + ) + + # Broadcast to all connected clients for this conversation + message_data = { + "type": "new_message", + "message": { + "id": str(message_id), + "conversation_id": conversation_id, + "role": "user", + "content": message_content, + "timestamp": datetime.now().isoformat(), + "status": "sent", + }, + } + + for conn in active_conversations.get(conversation_id, []): + try: + await conn.send_json(message_data) + except Exception: + # Connection might be closed; skip this subscriber + logger.debug("Skipped broadcast to a closed websocket") + + # TODO: Here we would trigger agent response generation + # For now, send a simple acknowledgment after a delay + await asyncio.sleep(1) + + agent_response = { + "type": "new_message", + "message": { + "id": str(uuid.uuid4()), + "conversation_id": conversation_id, + "role": "agent", + "content": f"I received your message: {message_content}", + "timestamp": datetime.now().isoformat(), + "status": "received", + }, + } + + for conn in active_conversations.get(conversation_id, []): + try: + await conn.send_json(agent_response) + except Exception: + # Connection might be closed; skip this subscriber + logger.debug("Skipped broadcast to a closed websocket") + + # Store agent response in database + async with get_db_connection() as conn: + await conn.execute( + """ + INSERT INTO agent_conversations (session_id, agent_id, message_type, content) + VALUES ($1, $2, 'agent', $3) + """, + conversation_id, + agent_id, + agent_response["message"]["content"], + ) + + except WebSocketDisconnect: + break + except Exception as e: + logger.error(f"Error in conversation WebSocket: {e}") + + except Exception as e: + logger.error(f"Conversation WebSocket error for {conversation_id}: {e}") + finally: + # Clean up connection + if ( + hasattr(app.state, "active_conversations") + and conversation_id in app.state.active_conversations + ): + if websocket in app.state.active_conversations[conversation_id]: + app.state.active_conversations[conversation_id].remove(websocket) + + +@app.websocket("/ws/organization/{organization_id}/updates") +async def websocket_organization_updates(websocket: WebSocket, organization_id: str): + """WebSocket endpoint for organization-wide updates""" + import uuid + + connection_id = f"org-{organization_id}-{uuid.uuid4()}" + + try: + connection = await websocket_manager.connect( + websocket=websocket, + connection_id=connection_id, + organization_id=organization_id, + subscriptions=[ + UpdateType.AGENT_CREATED.value, + UpdateType.AGENT_UPDATED.value, + UpdateType.AGENT_DELETED.value, + UpdateType.KNOWLEDGE_UPDATED.value, + UpdateType.KNOWLEDGE_INDEXED.value, + UpdateType.SYSTEM_NOTIFICATION.value, + ], + ) + + # Keep connection alive + while True: + try: + message = await websocket.receive_text() + if message == "ping": + await websocket.send_text("pong") + connection.last_ping = datetime.now() + except WebSocketDisconnect: + break + + except Exception as e: + logger.error(f"Organization WebSocket error for {organization_id}: {e}") + finally: + await websocket_manager.disconnect(connection_id) + + +# WebSocket Statistics Endpoint +@app.get("/ws/stats") +async def get_websocket_stats(): + """Get WebSocket connection statistics""" + try: + stats = websocket_manager.get_stats() + return stats + except Exception as e: + logger.error(f"Error getting WebSocket stats: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# Manual notification endpoints for testing +@app.post("/ws/test/agent/{agent_id}/status") +async def test_agent_status_notification( + agent_id: str, + status: str = Body(..., embed=True), + message: Optional[str] = Body(None, embed=True), +): + """Test endpoint to send agent status notifications""" + try: + await notify_agent_status_change( + agent_id=agent_id, + status=status, + additional_data={"message": message} if message else None, + ) + return {"status": "notification_sent", "agent_id": agent_id} + except Exception as e: + logger.error(f"Error sending test notification: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +# ============================================================================ +# Missing API Endpoints (Goals, Teams, Organizations) +# ============================================================================ + + +@app.get("/teams") +async def get_teams(): + """Get list of teams""" + # Mock data for now + return [ + { + "id": "1", + "name": "Development Team", + "description": "Frontend and backend developers", + "member_count": 5, + "organization_id": "1", + }, + { + "id": "2", + "name": "Executive Team", + "description": "Leadership and strategy", + "member_count": 3, + "organization_id": "1", + }, + ] + + +@app.get("/organizations/{organization_id}/goals") +async def get_organization_goals(organization_id: str): + """Get goals for an organization""" + # Mock data for now + return [ + { + "id": "1", + "title": "Increase Development Velocity", + "description": "Improve team productivity and code quality", + "status": "active", + "progress": 75, + "organization_id": organization_id, + "created_at": "2024-01-15T10:00:00Z", + "due_date": "2024-12-31T23:59:59Z", + }, + { + "id": "2", + "title": "Enhance AI Capabilities", + "description": "Expand AI agent capabilities and intelligence", + "status": "active", + "progress": 50, + "organization_id": organization_id, + "created_at": "2024-02-01T10:00:00Z", + "due_date": "2024-11-30T23:59:59Z", + }, + ] + + +@app.get("/goals/{goal_id}") +async def get_goal_details(goal_id: str): + """Get detailed information about a specific goal""" + # Mock data for now + return { + "id": goal_id, + "title": "Increase Development Velocity", + "description": "Improve team productivity and code quality through better tooling, processes, and automation", + "status": "active", + "progress": 75, + "organization_id": "1", + "team_id": "1", + "created_at": "2024-01-15T10:00:00Z", + "updated_at": "2024-08-06T16:30:00Z", + "due_date": "2024-12-31T23:59:59Z", + "milestones": [ + { + "id": "1", + "title": "Implement CI/CD Pipeline", + "description": "Set up automated testing and deployment", + "status": "completed", + "progress": 100, + "due_date": "2024-03-15T23:59:59Z", + }, + { + "id": "2", + "title": "Enhance Code Review Process", + "description": "Streamline code review workflow with automated tools", + "status": "in_progress", + "progress": 80, + "due_date": "2024-09-30T23:59:59Z", + }, + { + "id": "3", + "title": "Deploy AI-Powered Testing", + "description": "Implement intelligent test generation and execution", + "status": "planned", + "progress": 25, + "due_date": "2024-12-15T23:59:59Z", + }, + ], + "metrics": { + "deployment_frequency": "Daily", + "lead_time": "2.3 days", + "mttr": "45 minutes", + "change_failure_rate": "5%", + }, + "assigned_agents": [ + {"id": "1", "name": "DevOps Agent", "role": "CI/CD Specialist"}, + {"id": "2", "name": "QA Agent", "role": "Test Automation Engineer"}, + ], + } + + +@app.get("/agents/{agent_id}/tasks") +async def get_agent_tasks_list(agent_id: str): + """Get tasks for a specific agent - GET method""" + try: + # Get tasks from task queue + tasks = await app.state.task_queue.get_agent_tasks(agent_id) + return {"agent_id": agent_id, "tasks": tasks} + except Exception as e: + logger.error(f"Error getting tasks for agent {agent_id}: {e}") + raise HTTPException(status_code=500, detail=str(e)) diff --git a/services/orchestrator/mcp_integration.py b/services/orchestrator/mcp_integration.py index 8393511..34664a4 100644 --- a/services/orchestrator/mcp_integration.py +++ b/services/orchestrator/mcp_integration.py @@ -1,659 +1,659 @@ -""" -MCP (Model Context Protocol) Integration for FuzeAgent - -Provides MCP server functionality to give Claude SDK sessions access to: -- Organization structure and context -- Team information and agent hierarchy -- Agent capabilities and current status -- Task context and history -- Repository and project information - -This allows agents to have full organizational context when making decisions. -""" - -import asyncio -import json -import logging -import os -import uuid -from dataclasses import asdict, dataclass -from datetime import datetime -from typing import Any, Dict, List, Optional, Union - -from .database import DatabaseManager - -logger = logging.getLogger(__name__) - - -@dataclass -class MCPTool: - """Represents an MCP tool definition""" - - name: str - description: str - input_schema: Dict[str, Any] - - -@dataclass -class MCPResource: - """Represents an MCP resource""" - - uri: str - name: str - description: str - mime_type: str - - -class FuzeAgentMCPServer: - """ - MCP Server for FuzeAgent organizational context. - - Provides tools and resources for Claude SDK sessions to access: - - Organizational structure - - Agent capabilities and status - - Task context and history - - Repository information - """ - - def __init__(self): - self.tools = self._define_tools() - self.resources = self._define_resources() - - def _define_tools(self) -> List[MCPTool]: - """Define available MCP tools""" - return [ - MCPTool( - name="get_organization_structure", - description="Get the complete organizational structure including teams and agents", - input_schema={ - "type": "object", - "properties": { - "organization_id": { - "type": "string", - "description": "Optional organization ID to filter by", - } - }, - }, - ), - MCPTool( - name="get_team_agents", - description="Get all agents in a specific team with their capabilities", - input_schema={ - "type": "object", - "properties": { - "team_id": { - "type": "string", - "description": "Team ID to get agents for", - } - }, - "required": ["team_id"], - }, - ), - MCPTool( - name="get_agent_status", - description="Get current status and capabilities of a specific agent", - input_schema={ - "type": "object", - "properties": { - "agent_id": { - "type": "string", - "description": "Agent ID to get status for", - } - }, - "required": ["agent_id"], - }, - ), - MCPTool( - name="get_task_context", - description="Get comprehensive context for a task including history and related tasks", - input_schema={ - "type": "object", - "properties": { - "task_id": { - "type": "string", - "description": "Task ID to get context for", - }, - "include_history": { - "type": "boolean", - "description": "Whether to include task execution history", - "default": True, - }, - }, - "required": ["task_id"], - }, - ), - MCPTool( - name="get_agent_memory", - description="Get agent memory and previous interactions", - input_schema={ - "type": "object", - "properties": { - "agent_id": { - "type": "string", - "description": "Agent ID to get memory for", - }, - "limit": { - "type": "integer", - "description": "Maximum number of memory items to return", - "default": 10, - }, - "memory_type": { - "type": "string", - "description": "Type of memory to retrieve", - "enum": [ - "interactions", - "code_generations", - "performance_metrics", - ], - "default": "interactions", - }, - }, - "required": ["agent_id"], - }, - ), - MCPTool( - name="search_similar_tasks", - description="Search for similar tasks based on description or requirements", - input_schema={ - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "Search query for similar tasks", - }, - "agent_type": { - "type": "string", - "description": "Optional agent type to filter results", - }, - "limit": { - "type": "integer", - "description": "Maximum number of results", - "default": 5, - }, - }, - "required": ["query"], - }, - ), - MCPTool( - name="get_repository_context", - description="Get repository context and recent changes", - input_schema={ - "type": "object", - "properties": { - "repository_url": { - "type": "string", - "description": "Repository URL to get context for", - }, - "branch": { - "type": "string", - "description": "Optional branch name", - "default": "main", - }, - }, - "required": ["repository_url"], - }, - ), - MCPTool( - name="get_agent_recommendations", - description="Get agent recommendations for a specific task type", - input_schema={ - "type": "object", - "properties": { - "task_description": { - "type": "string", - "description": "Description of the task", - }, - "required_skills": { - "type": "array", - "items": {"type": "string"}, - "description": "List of required skills", - }, - "exclude_busy": { - "type": "boolean", - "description": "Whether to exclude currently busy agents", - "default": True, - }, - }, - "required": ["task_description"], - }, - ), - ] - - def _define_resources(self) -> List[MCPResource]: - """Define available MCP resources""" - return [ - MCPResource( - uri="fuzeagent://organizations", - name="Organizations", - description="Complete organizational structure and hierarchy", - mime_type="application/json", - ), - MCPResource( - uri="fuzeagent://agent-templates", - name="Agent Templates", - description="Available agent templates and their capabilities", - mime_type="application/json", - ), - MCPResource( - uri="fuzeagent://system-status", - name="System Status", - description="Current system status and health metrics", - mime_type="application/json", - ), - ] - - async def handle_tool_call( - self, tool_name: str, arguments: Dict[str, Any] - ) -> Dict[str, Any]: - """Handle MCP tool calls""" - try: - if tool_name == "get_organization_structure": - return await self._get_organization_structure(arguments) - elif tool_name == "get_team_agents": - return await self._get_team_agents(arguments) - elif tool_name == "get_agent_status": - return await self._get_agent_status(arguments) - elif tool_name == "get_task_context": - return await self._get_task_context(arguments) - elif tool_name == "get_agent_memory": - return await self._get_agent_memory(arguments) - elif tool_name == "search_similar_tasks": - return await self._search_similar_tasks(arguments) - elif tool_name == "get_repository_context": - return await self._get_repository_context(arguments) - elif tool_name == "get_agent_recommendations": - return await self._get_agent_recommendations(arguments) - else: - return {"error": f"Unknown tool: {tool_name}"} - - except Exception as e: - logger.error(f"Error in MCP tool call {tool_name}: {e}") - return {"error": str(e)} - - async def handle_resource_request(self, uri: str) -> Dict[str, Any]: - """Handle MCP resource requests""" - try: - if uri == "fuzeagent://organizations": - return await self._get_organizations_resource() - elif uri == "fuzeagent://agent-templates": - return await self._get_agent_templates_resource() - elif uri == "fuzeagent://system-status": - return await self._get_system_status_resource() - else: - return {"error": f"Unknown resource: {uri}"} - - except Exception as e: - logger.error(f"Error in MCP resource request {uri}: {e}") - return {"error": str(e)} - - # Tool implementations - - async def _get_organization_structure(self, args: Dict[str, Any]) -> Dict[str, Any]: - """Get organizational structure""" - organization_id = args.get("organization_id") - - # This would integrate with the MCP FuzeAgent server - # For now, return mock structure - return { - "organizations": [ - { - "id": "fuzeagent-org", - "name": "FuzeAgent Organization", - "teams": [ - { - "id": "dev-team-1", - "name": "Development Team Alpha", - "agents": [ - { - "id": "frontend-dev-1", - "name": "React Developer 1", - "type": "frontend_developer", - "status": "available", - "skills": ["react", "typescript", "css"], - }, - { - "id": "backend-dev-1", - "name": "Python Developer 1", - "type": "backend_developer", - "status": "busy", - "skills": ["python", "fastapi", "postgresql"], - }, - ], - } - ], - } - ] - } - - async def _get_team_agents(self, args: Dict[str, Any]) -> Dict[str, Any]: - """Get agents in a team""" - team_id = args["team_id"] - - # Get agents from database - agents = await DatabaseManager.get_agents_by_team(team_id) - - return { - "team_id": team_id, - "agents": [ - { - "id": agent["id"], - "name": agent["name"], - "role": agent["role"], - "type": agent["type"], - "status": agent["status"], - "capabilities": agent.get("config", {}).get("tools", []), - "current_task": agent.get("current_task_id"), - "created_at": ( - agent["created_at"].isoformat() if agent["created_at"] else None - ), - } - for agent in agents - ], - } - - async def _get_agent_status(self, args: Dict[str, Any]) -> Dict[str, Any]: - """Get agent status""" - agent_id = args["agent_id"] - - agent = await DatabaseManager.get_agent(agent_id) - if not agent: - return {"error": f"Agent {agent_id} not found"} - - # Get current tasks - tasks = await DatabaseManager.get_agent_tasks(agent_id, limit=5) - - return { - "agent_id": agent_id, - "name": agent["name"], - "role": agent["role"], - "type": agent["type"], - "status": agent["status"], - "capabilities": agent.get("config", {}).get("tools", []), - "model": agent.get("config", {}).get("model", "claude-sonnet-4-20250514"), - "current_tasks": [ - { - "id": task["id"], - "title": task["title"], - "status": task["status"], - "created_at": ( - task["created_at"].isoformat() if task["created_at"] else None - ), - } - for task in tasks - if task["status"] in ["pending", "executing"] - ], - "recent_tasks": [ - { - "id": task["id"], - "title": task["title"], - "status": task["status"], - "completed_at": ( - task["updated_at"].isoformat() if task["updated_at"] else None - ), - } - for task in tasks - if task["status"] in ["completed", "failed"] - ], - } - - async def _get_task_context(self, args: Dict[str, Any]) -> Dict[str, Any]: - """Get task context""" - task_id = args["task_id"] - include_history = args.get("include_history", True) - - # Get task data - task = await DatabaseManager.get_task(task_id) - if not task: - return {"error": f"Task {task_id} not found"} - - # Get agent data - agent = ( - await DatabaseManager.get_agent(task["assigned_to"]) - if task["assigned_to"] - else None - ) - - context = { - "task_id": task_id, - "title": task["title"], - "description": task["description"], - "status": task["status"], - "priority": task.get("priority", "medium"), - "created_at": ( - task["created_at"].isoformat() if task["created_at"] else None - ), - "assigned_agent": ( - { - "id": agent["id"], - "name": agent["name"], - "role": agent["role"], - "type": agent["type"], - } - if agent - else None - ), - } - - if include_history: - # Get task iterations - iterations = await DatabaseManager.get_task_iterations(task_id) - context["execution_history"] = [ - { - "iteration": iter["iteration_number"], - "step": iter["step"], - "started_at": ( - iter["started_at"].isoformat() if iter["started_at"] else None - ), - "completed_at": ( - iter["completed_at"].isoformat() - if iter["completed_at"] - else None - ), - "success": iter["success"], - "human_question": iter["human_question"], - "human_response": iter["human_response"], - } - for iter in iterations - ] - - return context - - async def _get_agent_memory(self, args: Dict[str, Any]) -> Dict[str, Any]: - """Get agent memory""" - agent_id = args["agent_id"] - limit = args.get("limit", 10) - memory_type = args.get("memory_type", "interactions") - - # This would integrate with conversation_manager - # For now return mock data - return { - "agent_id": agent_id, - "memory_type": memory_type, - "items": [ - { - "id": f"memory-{i}", - "type": memory_type, - "content": f"Sample {memory_type} {i}", - "timestamp": datetime.now().isoformat(), - "metadata": {}, - } - for i in range(min(limit, 5)) - ], - } - - async def _search_similar_tasks(self, args: Dict[str, Any]) -> Dict[str, Any]: - """Search for similar tasks""" - query = args["query"] - agent_type = args.get("agent_type") - limit = args.get("limit", 5) - - # This would use vector search in production - # For now return mock results - return { - "query": query, - "results": [ - { - "task_id": f"task-{i}", - "title": f"Similar task {i}", - "description": f"Task similar to '{query}'", - "similarity_score": 0.8 - (i * 0.1), - "agent_type": agent_type or "developer", - "status": "completed", - "completion_time_minutes": 120 + (i * 30), - } - for i in range(min(limit, 3)) - ], - } - - async def _get_repository_context(self, args: Dict[str, Any]) -> Dict[str, Any]: - """Get repository context""" - repository_url = args["repository_url"] - branch = args.get("branch", "main") - - return { - "repository_url": repository_url, - "branch": branch, - "recent_commits": [ - { - "hash": "abc123", - "message": "Recent commit message", - "author": "developer@example.com", - "timestamp": datetime.now().isoformat(), - } - ], - "active_branches": [branch, "develop", "feature/new-feature"], - "technologies": ["python", "fastapi", "react", "typescript"], - "structure": { - "backend": "services/orchestrator/", - "frontend": "services/ui-react/", - "containers": "containers/", - "docs": "docs/", - }, - } - - async def _get_agent_recommendations(self, args: Dict[str, Any]) -> Dict[str, Any]: - """Get agent recommendations for a task""" - task_description = args["task_description"] - required_skills = args.get("required_skills", []) - exclude_busy = args.get("exclude_busy", True) - - # This would use ML/AI to match agents to tasks - # For now return mock recommendations - return { - "task_description": task_description, - "required_skills": required_skills, - "recommendations": [ - { - "agent_id": "frontend-dev-1", - "name": "React Developer 1", - "match_score": 0.95, - "matching_skills": ["react", "typescript"], - "availability": "available", - "estimated_completion_time": "4-6 hours", - }, - { - "agent_id": "fullstack-dev-1", - "name": "Full Stack Developer 1", - "match_score": 0.85, - "matching_skills": ["react", "python"], - "availability": "busy_until_2pm", - "estimated_completion_time": "6-8 hours", - }, - ], - } - - # Resource implementations - - async def _get_organizations_resource(self) -> Dict[str, Any]: - """Get organizations resource""" - organizations = await DatabaseManager.get_organizations() - return {"content": organizations, "mime_type": "application/json"} - - async def _get_agent_templates_resource(self) -> Dict[str, Any]: - """Get agent templates resource""" - templates = await DatabaseManager.get_agent_templates() - return {"content": templates, "mime_type": "application/json"} - - async def _get_system_status_resource(self) -> Dict[str, Any]: - """Get system status resource""" - return { - "content": { - "status": "healthy", - "timestamp": datetime.now().isoformat(), - "active_agents": 5, - "running_tasks": 3, - "system_load": 0.45, - "memory_usage": 0.67, - }, - "mime_type": "application/json", - } - - -# MCP Server Integration with Claude SDK -class MCPClaudeIntegration: - """Integrates MCP server with Claude SDK sessions""" - - def __init__(self, mcp_server: FuzeAgentMCPServer): - self.mcp_server = mcp_server - - async def setup_claude_session_mcp( - self, session_id: str, agent_id: str, task_id: str - ) -> Dict[str, str]: - """Set up MCP tools for a Claude SDK session""" - - # Generate MCP server configuration for the session - mcp_config = { - "server_name": f"fuzeagent-{session_id}", - "server_command": ["python", "-m", "services.orchestrator.mcp_integration"], - "server_args": [ - "--session-id", - session_id, - "--agent-id", - agent_id, - "--task-id", - task_id, - ], - "environment": { - "FUZEAGENT_SESSION_ID": session_id, - "FUZEAGENT_AGENT_ID": agent_id, - "FUZEAGENT_TASK_ID": task_id, - }, - } - - # In production, this would configure Claude SDK to use this MCP server - # For now, return the configuration - return mcp_config - - async def get_session_context( - self, session_id: str, agent_id: str, task_id: str - ) -> Dict[str, Any]: - """Get comprehensive context for a Claude SDK session""" - - # Get agent context - agent_context = await self.mcp_server.handle_tool_call( - "get_agent_status", {"agent_id": agent_id} - ) - - # Get task context - task_context = await self.mcp_server.handle_tool_call( - "get_task_context", {"task_id": task_id} - ) - - # Get team context - if agent_context.get("team_id"): - team_context = await self.mcp_server.handle_tool_call( - "get_team_agents", {"team_id": agent_context["team_id"]} - ) - else: - team_context = {"agents": []} - - return { - "session_id": session_id, - "agent_context": agent_context, - "task_context": task_context, - "team_context": team_context, - "available_tools": [tool.name for tool in self.mcp_server.tools], - "available_resources": [ - resource.uri for resource in self.mcp_server.resources - ], - } +""" +MCP (Model Context Protocol) Integration for FuzeAgent + +Provides MCP server functionality to give Claude SDK sessions access to: +- Organization structure and context +- Team information and agent hierarchy +- Agent capabilities and current status +- Task context and history +- Repository and project information + +This allows agents to have full organizational context when making decisions. +""" + +import asyncio +import json +import logging +import os +import uuid +from dataclasses import asdict, dataclass +from datetime import datetime +from typing import Any, Dict, List, Optional, Union + +from .database import DatabaseManager + +logger = logging.getLogger(__name__) + + +@dataclass +class MCPTool: + """Represents an MCP tool definition""" + + name: str + description: str + input_schema: Dict[str, Any] + + +@dataclass +class MCPResource: + """Represents an MCP resource""" + + uri: str + name: str + description: str + mime_type: str + + +class FuzeAgentMCPServer: + """ + MCP Server for FuzeAgent organizational context. + + Provides tools and resources for Claude SDK sessions to access: + - Organizational structure + - Agent capabilities and status + - Task context and history + - Repository information + """ + + def __init__(self): + self.tools = self._define_tools() + self.resources = self._define_resources() + + def _define_tools(self) -> List[MCPTool]: + """Define available MCP tools""" + return [ + MCPTool( + name="get_organization_structure", + description="Get the complete organizational structure including teams and agents", + input_schema={ + "type": "object", + "properties": { + "organization_id": { + "type": "string", + "description": "Optional organization ID to filter by", + } + }, + }, + ), + MCPTool( + name="get_team_agents", + description="Get all agents in a specific team with their capabilities", + input_schema={ + "type": "object", + "properties": { + "team_id": { + "type": "string", + "description": "Team ID to get agents for", + } + }, + "required": ["team_id"], + }, + ), + MCPTool( + name="get_agent_status", + description="Get current status and capabilities of a specific agent", + input_schema={ + "type": "object", + "properties": { + "agent_id": { + "type": "string", + "description": "Agent ID to get status for", + } + }, + "required": ["agent_id"], + }, + ), + MCPTool( + name="get_task_context", + description="Get comprehensive context for a task including history and related tasks", + input_schema={ + "type": "object", + "properties": { + "task_id": { + "type": "string", + "description": "Task ID to get context for", + }, + "include_history": { + "type": "boolean", + "description": "Whether to include task execution history", + "default": True, + }, + }, + "required": ["task_id"], + }, + ), + MCPTool( + name="get_agent_memory", + description="Get agent memory and previous interactions", + input_schema={ + "type": "object", + "properties": { + "agent_id": { + "type": "string", + "description": "Agent ID to get memory for", + }, + "limit": { + "type": "integer", + "description": "Maximum number of memory items to return", + "default": 10, + }, + "memory_type": { + "type": "string", + "description": "Type of memory to retrieve", + "enum": [ + "interactions", + "code_generations", + "performance_metrics", + ], + "default": "interactions", + }, + }, + "required": ["agent_id"], + }, + ), + MCPTool( + name="search_similar_tasks", + description="Search for similar tasks based on description or requirements", + input_schema={ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Search query for similar tasks", + }, + "agent_type": { + "type": "string", + "description": "Optional agent type to filter results", + }, + "limit": { + "type": "integer", + "description": "Maximum number of results", + "default": 5, + }, + }, + "required": ["query"], + }, + ), + MCPTool( + name="get_repository_context", + description="Get repository context and recent changes", + input_schema={ + "type": "object", + "properties": { + "repository_url": { + "type": "string", + "description": "Repository URL to get context for", + }, + "branch": { + "type": "string", + "description": "Optional branch name", + "default": "main", + }, + }, + "required": ["repository_url"], + }, + ), + MCPTool( + name="get_agent_recommendations", + description="Get agent recommendations for a specific task type", + input_schema={ + "type": "object", + "properties": { + "task_description": { + "type": "string", + "description": "Description of the task", + }, + "required_skills": { + "type": "array", + "items": {"type": "string"}, + "description": "List of required skills", + }, + "exclude_busy": { + "type": "boolean", + "description": "Whether to exclude currently busy agents", + "default": True, + }, + }, + "required": ["task_description"], + }, + ), + ] + + def _define_resources(self) -> List[MCPResource]: + """Define available MCP resources""" + return [ + MCPResource( + uri="fuzeagent://organizations", + name="Organizations", + description="Complete organizational structure and hierarchy", + mime_type="application/json", + ), + MCPResource( + uri="fuzeagent://agent-templates", + name="Agent Templates", + description="Available agent templates and their capabilities", + mime_type="application/json", + ), + MCPResource( + uri="fuzeagent://system-status", + name="System Status", + description="Current system status and health metrics", + mime_type="application/json", + ), + ] + + async def handle_tool_call( + self, tool_name: str, arguments: Dict[str, Any] + ) -> Dict[str, Any]: + """Handle MCP tool calls""" + try: + if tool_name == "get_organization_structure": + return await self._get_organization_structure(arguments) + elif tool_name == "get_team_agents": + return await self._get_team_agents(arguments) + elif tool_name == "get_agent_status": + return await self._get_agent_status(arguments) + elif tool_name == "get_task_context": + return await self._get_task_context(arguments) + elif tool_name == "get_agent_memory": + return await self._get_agent_memory(arguments) + elif tool_name == "search_similar_tasks": + return await self._search_similar_tasks(arguments) + elif tool_name == "get_repository_context": + return await self._get_repository_context(arguments) + elif tool_name == "get_agent_recommendations": + return await self._get_agent_recommendations(arguments) + else: + return {"error": f"Unknown tool: {tool_name}"} + + except Exception as e: + logger.error(f"Error in MCP tool call {tool_name}: {e}") + return {"error": str(e)} + + async def handle_resource_request(self, uri: str) -> Dict[str, Any]: + """Handle MCP resource requests""" + try: + if uri == "fuzeagent://organizations": + return await self._get_organizations_resource() + elif uri == "fuzeagent://agent-templates": + return await self._get_agent_templates_resource() + elif uri == "fuzeagent://system-status": + return await self._get_system_status_resource() + else: + return {"error": f"Unknown resource: {uri}"} + + except Exception as e: + logger.error(f"Error in MCP resource request {uri}: {e}") + return {"error": str(e)} + + # Tool implementations + + async def _get_organization_structure(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Get organizational structure""" + organization_id = args.get("organization_id") + + # This would integrate with the MCP FuzeAgent server + # For now, return mock structure + return { + "organizations": [ + { + "id": "fuzeagent-org", + "name": "FuzeAgent Organization", + "teams": [ + { + "id": "dev-team-1", + "name": "Development Team Alpha", + "agents": [ + { + "id": "frontend-dev-1", + "name": "React Developer 1", + "type": "frontend_developer", + "status": "available", + "skills": ["react", "typescript", "css"], + }, + { + "id": "backend-dev-1", + "name": "Python Developer 1", + "type": "backend_developer", + "status": "busy", + "skills": ["python", "fastapi", "postgresql"], + }, + ], + } + ], + } + ] + } + + async def _get_team_agents(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Get agents in a team""" + team_id = args["team_id"] + + # Get agents from database + agents = await DatabaseManager.get_agents_by_team(team_id) + + return { + "team_id": team_id, + "agents": [ + { + "id": agent["id"], + "name": agent["name"], + "role": agent["role"], + "type": agent["type"], + "status": agent["status"], + "capabilities": agent.get("config", {}).get("tools", []), + "current_task": agent.get("current_task_id"), + "created_at": ( + agent["created_at"].isoformat() if agent["created_at"] else None + ), + } + for agent in agents + ], + } + + async def _get_agent_status(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Get agent status""" + agent_id = args["agent_id"] + + agent = await DatabaseManager.get_agent(agent_id) + if not agent: + return {"error": f"Agent {agent_id} not found"} + + # Get current tasks + tasks = await DatabaseManager.get_agent_tasks(agent_id, limit=5) + + return { + "agent_id": agent_id, + "name": agent["name"], + "role": agent["role"], + "type": agent["type"], + "status": agent["status"], + "capabilities": agent.get("config", {}).get("tools", []), + "model": agent.get("config", {}).get("model", "claude-sonnet-4-20250514"), + "current_tasks": [ + { + "id": task["id"], + "title": task["title"], + "status": task["status"], + "created_at": ( + task["created_at"].isoformat() if task["created_at"] else None + ), + } + for task in tasks + if task["status"] in ["pending", "executing"] + ], + "recent_tasks": [ + { + "id": task["id"], + "title": task["title"], + "status": task["status"], + "completed_at": ( + task["updated_at"].isoformat() if task["updated_at"] else None + ), + } + for task in tasks + if task["status"] in ["completed", "failed"] + ], + } + + async def _get_task_context(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Get task context""" + task_id = args["task_id"] + include_history = args.get("include_history", True) + + # Get task data + task = await DatabaseManager.get_task(task_id) + if not task: + return {"error": f"Task {task_id} not found"} + + # Get agent data + agent = ( + await DatabaseManager.get_agent(task["assigned_to"]) + if task["assigned_to"] + else None + ) + + context = { + "task_id": task_id, + "title": task["title"], + "description": task["description"], + "status": task["status"], + "priority": task.get("priority", "medium"), + "created_at": ( + task["created_at"].isoformat() if task["created_at"] else None + ), + "assigned_agent": ( + { + "id": agent["id"], + "name": agent["name"], + "role": agent["role"], + "type": agent["type"], + } + if agent + else None + ), + } + + if include_history: + # Get task iterations + iterations = await DatabaseManager.get_task_iterations(task_id) + context["execution_history"] = [ + { + "iteration": iter["iteration_number"], + "step": iter["step"], + "started_at": ( + iter["started_at"].isoformat() if iter["started_at"] else None + ), + "completed_at": ( + iter["completed_at"].isoformat() + if iter["completed_at"] + else None + ), + "success": iter["success"], + "human_question": iter["human_question"], + "human_response": iter["human_response"], + } + for iter in iterations + ] + + return context + + async def _get_agent_memory(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Get agent memory""" + agent_id = args["agent_id"] + limit = args.get("limit", 10) + memory_type = args.get("memory_type", "interactions") + + # This would integrate with conversation_manager + # For now return mock data + return { + "agent_id": agent_id, + "memory_type": memory_type, + "items": [ + { + "id": f"memory-{i}", + "type": memory_type, + "content": f"Sample {memory_type} {i}", + "timestamp": datetime.now().isoformat(), + "metadata": {}, + } + for i in range(min(limit, 5)) + ], + } + + async def _search_similar_tasks(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Search for similar tasks""" + query = args["query"] + agent_type = args.get("agent_type") + limit = args.get("limit", 5) + + # This would use vector search in production + # For now return mock results + return { + "query": query, + "results": [ + { + "task_id": f"task-{i}", + "title": f"Similar task {i}", + "description": f"Task similar to '{query}'", + "similarity_score": 0.8 - (i * 0.1), + "agent_type": agent_type or "developer", + "status": "completed", + "completion_time_minutes": 120 + (i * 30), + } + for i in range(min(limit, 3)) + ], + } + + async def _get_repository_context(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Get repository context""" + repository_url = args["repository_url"] + branch = args.get("branch", "main") + + return { + "repository_url": repository_url, + "branch": branch, + "recent_commits": [ + { + "hash": "abc123", + "message": "Recent commit message", + "author": "developer@example.com", + "timestamp": datetime.now().isoformat(), + } + ], + "active_branches": [branch, "develop", "feature/new-feature"], + "technologies": ["python", "fastapi", "react", "typescript"], + "structure": { + "backend": "services/orchestrator/", + "frontend": "services/ui-react/", + "containers": "containers/", + "docs": "docs/", + }, + } + + async def _get_agent_recommendations(self, args: Dict[str, Any]) -> Dict[str, Any]: + """Get agent recommendations for a task""" + task_description = args["task_description"] + required_skills = args.get("required_skills", []) + exclude_busy = args.get("exclude_busy", True) + + # This would use ML/AI to match agents to tasks + # For now return mock recommendations + return { + "task_description": task_description, + "required_skills": required_skills, + "recommendations": [ + { + "agent_id": "frontend-dev-1", + "name": "React Developer 1", + "match_score": 0.95, + "matching_skills": ["react", "typescript"], + "availability": "available", + "estimated_completion_time": "4-6 hours", + }, + { + "agent_id": "fullstack-dev-1", + "name": "Full Stack Developer 1", + "match_score": 0.85, + "matching_skills": ["react", "python"], + "availability": "busy_until_2pm", + "estimated_completion_time": "6-8 hours", + }, + ], + } + + # Resource implementations + + async def _get_organizations_resource(self) -> Dict[str, Any]: + """Get organizations resource""" + organizations = await DatabaseManager.get_organizations() + return {"content": organizations, "mime_type": "application/json"} + + async def _get_agent_templates_resource(self) -> Dict[str, Any]: + """Get agent templates resource""" + templates = await DatabaseManager.get_agent_templates() + return {"content": templates, "mime_type": "application/json"} + + async def _get_system_status_resource(self) -> Dict[str, Any]: + """Get system status resource""" + return { + "content": { + "status": "healthy", + "timestamp": datetime.now().isoformat(), + "active_agents": 5, + "running_tasks": 3, + "system_load": 0.45, + "memory_usage": 0.67, + }, + "mime_type": "application/json", + } + + +# MCP Server Integration with Claude SDK +class MCPClaudeIntegration: + """Integrates MCP server with Claude SDK sessions""" + + def __init__(self, mcp_server: FuzeAgentMCPServer): + self.mcp_server = mcp_server + + async def setup_claude_session_mcp( + self, session_id: str, agent_id: str, task_id: str + ) -> Dict[str, str]: + """Set up MCP tools for a Claude SDK session""" + + # Generate MCP server configuration for the session + mcp_config = { + "server_name": f"fuzeagent-{session_id}", + "server_command": ["python", "-m", "services.orchestrator.mcp_integration"], + "server_args": [ + "--session-id", + session_id, + "--agent-id", + agent_id, + "--task-id", + task_id, + ], + "environment": { + "FUZEAGENT_SESSION_ID": session_id, + "FUZEAGENT_AGENT_ID": agent_id, + "FUZEAGENT_TASK_ID": task_id, + }, + } + + # In production, this would configure Claude SDK to use this MCP server + # For now, return the configuration + return mcp_config + + async def get_session_context( + self, session_id: str, agent_id: str, task_id: str + ) -> Dict[str, Any]: + """Get comprehensive context for a Claude SDK session""" + + # Get agent context + agent_context = await self.mcp_server.handle_tool_call( + "get_agent_status", {"agent_id": agent_id} + ) + + # Get task context + task_context = await self.mcp_server.handle_tool_call( + "get_task_context", {"task_id": task_id} + ) + + # Get team context + if agent_context.get("team_id"): + team_context = await self.mcp_server.handle_tool_call( + "get_team_agents", {"team_id": agent_context["team_id"]} + ) + else: + team_context = {"agents": []} + + return { + "session_id": session_id, + "agent_context": agent_context, + "task_context": task_context, + "team_context": team_context, + "available_tools": [tool.name for tool in self.mcp_server.tools], + "available_resources": [ + resource.uri for resource in self.mcp_server.resources + ], + } diff --git a/services/orchestrator/model_configuration.py b/services/orchestrator/model_configuration.py index e0834b1..de5cf16 100644 --- a/services/orchestrator/model_configuration.py +++ b/services/orchestrator/model_configuration.py @@ -1,558 +1,558 @@ -""" -Model Configuration and API Key Management for FuzeAgent - -Manages model configurations and API keys for different AI providers at the -organization level, with secure storage and agent-specific model selection. - -Supports: -- Multiple AI model providers (Anthropic, OpenAI, Google, etc.) -- Organization-level API key management -- Agent-specific model configurations -- Secure credential storage and access -- Model capability matching -- Cost optimization -""" - -import base64 -import json -import logging -import os -import tempfile -from dataclasses import dataclass, field -from datetime import datetime -from enum import Enum -from typing import Any, Dict, List, Optional, Union - -from cryptography.fernet import Fernet - -from .database import DatabaseManager - -logger = logging.getLogger(__name__) - - -class ModelProvider(str, Enum): - ANTHROPIC = "anthropic" - OPENAI = "openai" - GOOGLE = "google" - AZURE_OPENAI = "azure_openai" - COHERE = "cohere" - HUGGINGFACE = "huggingface" - OLLAMA = "ollama" - CUSTOM = "custom" - - -class ModelCapability(str, Enum): - TEXT_GENERATION = "text_generation" - CODE_GENERATION = "code_generation" - REASONING = "reasoning" - ANALYSIS = "analysis" - CONVERSATION = "conversation" - FUNCTION_CALLING = "function_calling" - VISION = "vision" - EMBEDDINGS = "embeddings" - - -@dataclass -class ModelSpec: - """Specification for an AI model""" - - model_id: str - provider: ModelProvider - name: str - description: str - capabilities: List[ModelCapability] - context_window: int - max_output_tokens: int - cost_per_input_token: float # USD per 1K tokens - cost_per_output_token: float # USD per 1K tokens - supports_streaming: bool = True - supports_function_calling: bool = False - supports_vision: bool = False - supports_json_mode: bool = False - temperature_range: tuple = (0.0, 2.0) - recommended_use_cases: List[str] = field(default_factory=list) - - -@dataclass -class ProviderCredentials: - """Secure storage for provider API credentials""" - - provider: ModelProvider - encrypted_api_key: str - endpoint_url: Optional[str] = None - additional_config: Dict[str, Any] = field(default_factory=dict) - created_at: datetime = field(default_factory=datetime.now) - last_used: Optional[datetime] = None - is_active: bool = True - - -@dataclass -class AgentModelConfig: - """Model configuration for an agent""" - - agent_id: str - primary_model: str # model_id - fallback_models: List[str] = field(default_factory=list) - temperature: float = 0.7 - max_tokens: int = 4096 - top_p: float = 1.0 - frequency_penalty: float = 0.0 - presence_penalty: float = 0.0 - custom_instructions: str = "" - use_function_calling: bool = True - streaming_enabled: bool = True - cost_limit_per_task: Optional[float] = None # USD - created_at: datetime = field(default_factory=datetime.now) - updated_at: datetime = field(default_factory=datetime.now) - - -class ModelConfigurationManager: - """ - Manages AI model configurations and provider credentials. - - Features: - - Secure API key storage with encryption - - Organization-level credential management - - Agent-specific model configurations - - Model capability matching - - Cost tracking and limits - - Automatic fallback handling - """ - - def __init__(self): - self.encryption_key = self._get_or_create_encryption_key() - self.fernet = Fernet(self.encryption_key) - self.available_models = self._initialize_model_specs() - - def _get_or_create_encryption_key(self) -> bytes: - """Get or create encryption key for API credentials""" - # Avoid a hard-coded world-readable /tmp path; allow override and fall - # back to the platform temp dir (still /tmp inside the Linux container). - key_dir = os.getenv("FUZEAGENT_KEY_DIR", tempfile.gettempdir()) - key_file = os.path.join(key_dir, "fuzeagent_encryption.key") - - if os.path.exists(key_file): - with open(key_file, "rb") as f: - return f.read() - else: - key = Fernet.generate_key() - with open(key_file, "wb") as f: - f.write(key) - return key - - def _initialize_model_specs(self) -> Dict[str, ModelSpec]: - """Initialize available model specifications""" - models = {} - - # Anthropic Claude models - models["claude-3-5-sonnet-20241022"] = ModelSpec( - model_id="claude-3-5-sonnet-20241022", - provider=ModelProvider.ANTHROPIC, - name="Claude 3.5 Sonnet", - description="Most intelligent model for complex reasoning and coding", - capabilities=[ - ModelCapability.TEXT_GENERATION, - ModelCapability.CODE_GENERATION, - ModelCapability.REASONING, - ModelCapability.ANALYSIS, - ModelCapability.CONVERSATION, - ModelCapability.FUNCTION_CALLING, - ModelCapability.VISION, - ], - context_window=200000, - max_output_tokens=8192, - cost_per_input_token=0.003, - cost_per_output_token=0.015, - supports_function_calling=True, - supports_vision=True, - supports_json_mode=True, - recommended_use_cases=[ - "complex coding", - "reasoning", - "analysis", - "research", - ], - ) - - models["claude-3-haiku-20240307"] = ModelSpec( - model_id="claude-3-haiku-20240307", - provider=ModelProvider.ANTHROPIC, - name="Claude 3 Haiku", - description="Fastest and most cost-effective model for simple tasks", - capabilities=[ - ModelCapability.TEXT_GENERATION, - ModelCapability.CODE_GENERATION, - ModelCapability.CONVERSATION, - ], - context_window=200000, - max_output_tokens=4096, - cost_per_input_token=0.00025, - cost_per_output_token=0.00125, - supports_function_calling=True, - supports_vision=True, - recommended_use_cases=[ - "simple tasks", - "quick responses", - "cost optimization", - ], - ) - - # OpenAI GPT models - models["gpt-4o"] = ModelSpec( - model_id="gpt-4o", - provider=ModelProvider.OPENAI, - name="GPT-4 Omni", - description="OpenAI's most capable multimodal model", - capabilities=[ - ModelCapability.TEXT_GENERATION, - ModelCapability.CODE_GENERATION, - ModelCapability.REASONING, - ModelCapability.ANALYSIS, - ModelCapability.CONVERSATION, - ModelCapability.FUNCTION_CALLING, - ModelCapability.VISION, - ], - context_window=128000, - max_output_tokens=4096, - cost_per_input_token=0.005, - cost_per_output_token=0.015, - supports_function_calling=True, - supports_vision=True, - supports_json_mode=True, - recommended_use_cases=[ - "multimodal tasks", - "function calling", - "complex reasoning", - ], - ) - - models["gpt-4o-mini"] = ModelSpec( - model_id="gpt-4o-mini", - provider=ModelProvider.OPENAI, - name="GPT-4 Omni Mini", - description="Cost-effective model for simpler tasks", - capabilities=[ - ModelCapability.TEXT_GENERATION, - ModelCapability.CODE_GENERATION, - ModelCapability.CONVERSATION, - ModelCapability.FUNCTION_CALLING, - ], - context_window=128000, - max_output_tokens=16384, - cost_per_input_token=0.00015, - cost_per_output_token=0.0006, - supports_function_calling=True, - supports_json_mode=True, - recommended_use_cases=["cost optimization", "simple tasks", "high volume"], - ) - - # Google Gemini models - models["gemini-1.5-pro"] = ModelSpec( - model_id="gemini-1.5-pro", - provider=ModelProvider.GOOGLE, - name="Gemini 1.5 Pro", - description="Google's most capable model with long context", - capabilities=[ - ModelCapability.TEXT_GENERATION, - ModelCapability.CODE_GENERATION, - ModelCapability.REASONING, - ModelCapability.ANALYSIS, - ModelCapability.CONVERSATION, - ModelCapability.FUNCTION_CALLING, - ModelCapability.VISION, - ], - context_window=2000000, # 2M tokens - max_output_tokens=8192, - cost_per_input_token=0.00125, - cost_per_output_token=0.005, - supports_function_calling=True, - supports_vision=True, - recommended_use_cases=[ - "long context", - "document analysis", - "multimodal tasks", - ], - ) - - return models - - async def store_provider_credentials( - self, - organization_id: str, - provider: ModelProvider, - api_key: str, - endpoint_url: Optional[str] = None, - additional_config: Optional[Dict[str, Any]] = None, - ) -> bool: - """Store encrypted API credentials for a provider""" - try: - # Encrypt the API key - encrypted_key = self.fernet.encrypt(api_key.encode()).decode() - - credentials = ProviderCredentials( - provider=provider, - encrypted_api_key=encrypted_key, - endpoint_url=endpoint_url, - additional_config=additional_config or {}, - ) - - # Store in database - await DatabaseManager.store_provider_credentials( - organization_id, credentials.__dict__ - ) - - logger.info( - f"Stored credentials for {provider} in organization {organization_id}" - ) - return True - - except Exception as e: - logger.error(f"Error storing provider credentials: {e}") - return False - - async def get_provider_credentials( - self, organization_id: str, provider: ModelProvider - ) -> Optional[str]: - """Get decrypted API key for a provider""" - try: - credentials_data = await DatabaseManager.get_provider_credentials( - organization_id, provider.value - ) - - if not credentials_data: - return None - - # Decrypt the API key - encrypted_key = credentials_data["encrypted_api_key"] - decrypted_key = self.fernet.decrypt(encrypted_key.encode()).decode() - - # Update last used timestamp - await DatabaseManager.update_credentials_last_used( - organization_id, provider.value - ) - - return decrypted_key - - except Exception as e: - logger.error(f"Error retrieving provider credentials: {e}") - return None - - async def configure_agent_model( - self, agent_id: str, model_config: AgentModelConfig - ) -> bool: - """Configure model settings for an agent""" - try: - # Validate primary model exists - if model_config.primary_model not in self.available_models: - raise ValueError(f"Unknown model: {model_config.primary_model}") - - # Validate fallback models - for model_id in model_config.fallback_models: - if model_id not in self.available_models: - raise ValueError(f"Unknown fallback model: {model_id}") - - # Store configuration - await DatabaseManager.store_agent_model_config( - agent_id, model_config.__dict__ - ) - - logger.info(f"Configured model settings for agent {agent_id}") - return True - - except Exception as e: - logger.error(f"Error configuring agent model: {e}") - return False - - async def get_agent_model_config(self, agent_id: str) -> Optional[AgentModelConfig]: - """Get model configuration for an agent""" - try: - config_data = await DatabaseManager.get_agent_model_config(agent_id) - - if not config_data: - # Return default configuration - return AgentModelConfig( - agent_id=agent_id, - primary_model="claude-3-5-sonnet-20241022", # Default to Claude 3.5 Sonnet - ) - - return AgentModelConfig(**config_data) - - except Exception as e: - logger.error(f"Error getting agent model config: {e}") - return None - - async def get_model_for_task( - self, - agent_id: str, - task_capabilities: List[ModelCapability], - cost_limit: Optional[float] = None, - ) -> Optional[str]: - """Select best model for a task based on capabilities and cost""" - try: - agent_config = await self.get_agent_model_config(agent_id) - if not agent_config: - return None - - # Check if primary model supports required capabilities - primary_model = self.available_models.get(agent_config.primary_model) - if primary_model and all( - cap in primary_model.capabilities for cap in task_capabilities - ): - # Check cost limit if specified - if ( - cost_limit is None - or self._estimate_task_cost(primary_model, 1000) <= cost_limit - ): - return agent_config.primary_model - - # Try fallback models - for model_id in agent_config.fallback_models: - model = self.available_models.get(model_id) - if model and all( - cap in model.capabilities for cap in task_capabilities - ): - if ( - cost_limit is None - or self._estimate_task_cost(model, 1000) <= cost_limit - ): - return model_id - - # No suitable model found - logger.warning( - f"No suitable model found for agent {agent_id} with capabilities {task_capabilities}" - ) - return None - - except Exception as e: - logger.error(f"Error selecting model for task: {e}") - return None - - def _estimate_task_cost(self, model: ModelSpec, estimated_tokens: int) -> float: - """Estimate cost for a task with given token count""" - # Simple estimation assuming 70% input, 30% output tokens - input_tokens = int(estimated_tokens * 0.7) - output_tokens = int(estimated_tokens * 0.3) - - input_cost = (input_tokens / 1000) * model.cost_per_input_token - output_cost = (output_tokens / 1000) * model.cost_per_output_token - - return input_cost + output_cost - - async def get_available_models( - self, - organization_id: str, - provider: Optional[ModelProvider] = None, - capabilities: Optional[List[ModelCapability]] = None, - ) -> List[Dict[str, Any]]: - """Get available models with provider credential validation""" - available = [] - - for model_id, model in self.available_models.items(): - # Filter by provider if specified - if provider and model.provider != provider: - continue - - # Filter by capabilities if specified - if capabilities and not all( - cap in model.capabilities for cap in capabilities - ): - continue - - # Check if organization has credentials for this provider - has_credentials = ( - await self.get_provider_credentials(organization_id, model.provider) - is not None - ) - - model_info = { - "model_id": model_id, - "provider": model.provider.value, - "name": model.name, - "description": model.description, - "capabilities": [cap.value for cap in model.capabilities], - "context_window": model.context_window, - "max_output_tokens": model.max_output_tokens, - "cost_per_input_token": model.cost_per_input_token, - "cost_per_output_token": model.cost_per_output_token, - "supports_streaming": model.supports_streaming, - "supports_function_calling": model.supports_function_calling, - "supports_vision": model.supports_vision, - "supports_json_mode": model.supports_json_mode, - "recommended_use_cases": model.recommended_use_cases, - "has_credentials": has_credentials, - "available": has_credentials, - } - - available.append(model_info) - - return available - - async def get_organization_model_usage( - self, organization_id: str, days: int = 30 - ) -> Dict[str, Any]: - """Get model usage statistics for an organization""" - try: - usage_data = await DatabaseManager.get_model_usage_stats( - organization_id, days - ) - - return { - "organization_id": organization_id, - "period_days": days, - "total_requests": usage_data.get("total_requests", 0), - "total_tokens": usage_data.get("total_tokens", 0), - "total_cost": usage_data.get("total_cost", 0.0), - "model_breakdown": usage_data.get("model_breakdown", {}), - "agent_breakdown": usage_data.get("agent_breakdown", {}), - "daily_usage": usage_data.get("daily_usage", []), - } - - except Exception as e: - logger.error(f"Error getting model usage: {e}") - return {} - - async def estimate_task_cost( - self, agent_id: str, task_description: str, estimated_complexity: str = "medium" - ) -> Dict[str, Any]: - """Estimate cost for a task execution""" - try: - agent_config = await self.get_agent_model_config(agent_id) - if not agent_config: - return {"error": "Agent configuration not found"} - - model = self.available_models.get(agent_config.primary_model) - if not model: - return {"error": "Model specification not found"} - - # Estimate token usage based on complexity - token_estimates = { - "low": 2000, - "medium": 5000, - "high": 10000, - "very_high": 20000, - } - - estimated_tokens = token_estimates.get(estimated_complexity, 5000) - estimated_cost = self._estimate_task_cost(model, estimated_tokens) - - return { - "agent_id": agent_id, - "model": model.model_id, - "estimated_tokens": estimated_tokens, - "estimated_cost_usd": round(estimated_cost, 4), - "complexity": estimated_complexity, - "cost_breakdown": { - "input_cost": (estimated_tokens * 0.7 / 1000) - * model.cost_per_input_token, - "output_cost": (estimated_tokens * 0.3 / 1000) - * model.cost_per_output_token, - }, - } - - except Exception as e: - logger.error(f"Error estimating task cost: {e}") - return {"error": str(e)} - - -# Global instance -model_config_manager = ModelConfigurationManager() +""" +Model Configuration and API Key Management for FuzeAgent + +Manages model configurations and API keys for different AI providers at the +organization level, with secure storage and agent-specific model selection. + +Supports: +- Multiple AI model providers (Anthropic, OpenAI, Google, etc.) +- Organization-level API key management +- Agent-specific model configurations +- Secure credential storage and access +- Model capability matching +- Cost optimization +""" + +import base64 +import json +import logging +import os +import tempfile +from dataclasses import dataclass, field +from datetime import datetime +from enum import Enum +from typing import Any, Dict, List, Optional, Union + +from cryptography.fernet import Fernet + +from .database import DatabaseManager + +logger = logging.getLogger(__name__) + + +class ModelProvider(str, Enum): + ANTHROPIC = "anthropic" + OPENAI = "openai" + GOOGLE = "google" + AZURE_OPENAI = "azure_openai" + COHERE = "cohere" + HUGGINGFACE = "huggingface" + OLLAMA = "ollama" + CUSTOM = "custom" + + +class ModelCapability(str, Enum): + TEXT_GENERATION = "text_generation" + CODE_GENERATION = "code_generation" + REASONING = "reasoning" + ANALYSIS = "analysis" + CONVERSATION = "conversation" + FUNCTION_CALLING = "function_calling" + VISION = "vision" + EMBEDDINGS = "embeddings" + + +@dataclass +class ModelSpec: + """Specification for an AI model""" + + model_id: str + provider: ModelProvider + name: str + description: str + capabilities: List[ModelCapability] + context_window: int + max_output_tokens: int + cost_per_input_token: float # USD per 1K tokens + cost_per_output_token: float # USD per 1K tokens + supports_streaming: bool = True + supports_function_calling: bool = False + supports_vision: bool = False + supports_json_mode: bool = False + temperature_range: tuple = (0.0, 2.0) + recommended_use_cases: List[str] = field(default_factory=list) + + +@dataclass +class ProviderCredentials: + """Secure storage for provider API credentials""" + + provider: ModelProvider + encrypted_api_key: str + endpoint_url: Optional[str] = None + additional_config: Dict[str, Any] = field(default_factory=dict) + created_at: datetime = field(default_factory=datetime.now) + last_used: Optional[datetime] = None + is_active: bool = True + + +@dataclass +class AgentModelConfig: + """Model configuration for an agent""" + + agent_id: str + primary_model: str # model_id + fallback_models: List[str] = field(default_factory=list) + temperature: float = 0.7 + max_tokens: int = 4096 + top_p: float = 1.0 + frequency_penalty: float = 0.0 + presence_penalty: float = 0.0 + custom_instructions: str = "" + use_function_calling: bool = True + streaming_enabled: bool = True + cost_limit_per_task: Optional[float] = None # USD + created_at: datetime = field(default_factory=datetime.now) + updated_at: datetime = field(default_factory=datetime.now) + + +class ModelConfigurationManager: + """ + Manages AI model configurations and provider credentials. + + Features: + - Secure API key storage with encryption + - Organization-level credential management + - Agent-specific model configurations + - Model capability matching + - Cost tracking and limits + - Automatic fallback handling + """ + + def __init__(self): + self.encryption_key = self._get_or_create_encryption_key() + self.fernet = Fernet(self.encryption_key) + self.available_models = self._initialize_model_specs() + + def _get_or_create_encryption_key(self) -> bytes: + """Get or create encryption key for API credentials""" + # Avoid a hard-coded world-readable /tmp path; allow override and fall + # back to the platform temp dir (still /tmp inside the Linux container). + key_dir = os.getenv("FUZEAGENT_KEY_DIR", tempfile.gettempdir()) + key_file = os.path.join(key_dir, "fuzeagent_encryption.key") + + if os.path.exists(key_file): + with open(key_file, "rb") as f: + return f.read() + else: + key = Fernet.generate_key() + with open(key_file, "wb") as f: + f.write(key) + return key + + def _initialize_model_specs(self) -> Dict[str, ModelSpec]: + """Initialize available model specifications""" + models = {} + + # Anthropic Claude models + models["claude-3-5-sonnet-20241022"] = ModelSpec( + model_id="claude-3-5-sonnet-20241022", + provider=ModelProvider.ANTHROPIC, + name="Claude 3.5 Sonnet", + description="Most intelligent model for complex reasoning and coding", + capabilities=[ + ModelCapability.TEXT_GENERATION, + ModelCapability.CODE_GENERATION, + ModelCapability.REASONING, + ModelCapability.ANALYSIS, + ModelCapability.CONVERSATION, + ModelCapability.FUNCTION_CALLING, + ModelCapability.VISION, + ], + context_window=200000, + max_output_tokens=8192, + cost_per_input_token=0.003, + cost_per_output_token=0.015, + supports_function_calling=True, + supports_vision=True, + supports_json_mode=True, + recommended_use_cases=[ + "complex coding", + "reasoning", + "analysis", + "research", + ], + ) + + models["claude-3-haiku-20240307"] = ModelSpec( + model_id="claude-3-haiku-20240307", + provider=ModelProvider.ANTHROPIC, + name="Claude 3 Haiku", + description="Fastest and most cost-effective model for simple tasks", + capabilities=[ + ModelCapability.TEXT_GENERATION, + ModelCapability.CODE_GENERATION, + ModelCapability.CONVERSATION, + ], + context_window=200000, + max_output_tokens=4096, + cost_per_input_token=0.00025, + cost_per_output_token=0.00125, + supports_function_calling=True, + supports_vision=True, + recommended_use_cases=[ + "simple tasks", + "quick responses", + "cost optimization", + ], + ) + + # OpenAI GPT models + models["gpt-4o"] = ModelSpec( + model_id="gpt-4o", + provider=ModelProvider.OPENAI, + name="GPT-4 Omni", + description="OpenAI's most capable multimodal model", + capabilities=[ + ModelCapability.TEXT_GENERATION, + ModelCapability.CODE_GENERATION, + ModelCapability.REASONING, + ModelCapability.ANALYSIS, + ModelCapability.CONVERSATION, + ModelCapability.FUNCTION_CALLING, + ModelCapability.VISION, + ], + context_window=128000, + max_output_tokens=4096, + cost_per_input_token=0.005, + cost_per_output_token=0.015, + supports_function_calling=True, + supports_vision=True, + supports_json_mode=True, + recommended_use_cases=[ + "multimodal tasks", + "function calling", + "complex reasoning", + ], + ) + + models["gpt-4o-mini"] = ModelSpec( + model_id="gpt-4o-mini", + provider=ModelProvider.OPENAI, + name="GPT-4 Omni Mini", + description="Cost-effective model for simpler tasks", + capabilities=[ + ModelCapability.TEXT_GENERATION, + ModelCapability.CODE_GENERATION, + ModelCapability.CONVERSATION, + ModelCapability.FUNCTION_CALLING, + ], + context_window=128000, + max_output_tokens=16384, + cost_per_input_token=0.00015, + cost_per_output_token=0.0006, + supports_function_calling=True, + supports_json_mode=True, + recommended_use_cases=["cost optimization", "simple tasks", "high volume"], + ) + + # Google Gemini models + models["gemini-1.5-pro"] = ModelSpec( + model_id="gemini-1.5-pro", + provider=ModelProvider.GOOGLE, + name="Gemini 1.5 Pro", + description="Google's most capable model with long context", + capabilities=[ + ModelCapability.TEXT_GENERATION, + ModelCapability.CODE_GENERATION, + ModelCapability.REASONING, + ModelCapability.ANALYSIS, + ModelCapability.CONVERSATION, + ModelCapability.FUNCTION_CALLING, + ModelCapability.VISION, + ], + context_window=2000000, # 2M tokens + max_output_tokens=8192, + cost_per_input_token=0.00125, + cost_per_output_token=0.005, + supports_function_calling=True, + supports_vision=True, + recommended_use_cases=[ + "long context", + "document analysis", + "multimodal tasks", + ], + ) + + return models + + async def store_provider_credentials( + self, + organization_id: str, + provider: ModelProvider, + api_key: str, + endpoint_url: Optional[str] = None, + additional_config: Optional[Dict[str, Any]] = None, + ) -> bool: + """Store encrypted API credentials for a provider""" + try: + # Encrypt the API key + encrypted_key = self.fernet.encrypt(api_key.encode()).decode() + + credentials = ProviderCredentials( + provider=provider, + encrypted_api_key=encrypted_key, + endpoint_url=endpoint_url, + additional_config=additional_config or {}, + ) + + # Store in database + await DatabaseManager.store_provider_credentials( + organization_id, credentials.__dict__ + ) + + logger.info( + f"Stored credentials for {provider} in organization {organization_id}" + ) + return True + + except Exception as e: + logger.error(f"Error storing provider credentials: {e}") + return False + + async def get_provider_credentials( + self, organization_id: str, provider: ModelProvider + ) -> Optional[str]: + """Get decrypted API key for a provider""" + try: + credentials_data = await DatabaseManager.get_provider_credentials( + organization_id, provider.value + ) + + if not credentials_data: + return None + + # Decrypt the API key + encrypted_key = credentials_data["encrypted_api_key"] + decrypted_key = self.fernet.decrypt(encrypted_key.encode()).decode() + + # Update last used timestamp + await DatabaseManager.update_credentials_last_used( + organization_id, provider.value + ) + + return decrypted_key + + except Exception as e: + logger.error(f"Error retrieving provider credentials: {e}") + return None + + async def configure_agent_model( + self, agent_id: str, model_config: AgentModelConfig + ) -> bool: + """Configure model settings for an agent""" + try: + # Validate primary model exists + if model_config.primary_model not in self.available_models: + raise ValueError(f"Unknown model: {model_config.primary_model}") + + # Validate fallback models + for model_id in model_config.fallback_models: + if model_id not in self.available_models: + raise ValueError(f"Unknown fallback model: {model_id}") + + # Store configuration + await DatabaseManager.store_agent_model_config( + agent_id, model_config.__dict__ + ) + + logger.info(f"Configured model settings for agent {agent_id}") + return True + + except Exception as e: + logger.error(f"Error configuring agent model: {e}") + return False + + async def get_agent_model_config(self, agent_id: str) -> Optional[AgentModelConfig]: + """Get model configuration for an agent""" + try: + config_data = await DatabaseManager.get_agent_model_config(agent_id) + + if not config_data: + # Return default configuration + return AgentModelConfig( + agent_id=agent_id, + primary_model="claude-3-5-sonnet-20241022", # Default to Claude 3.5 Sonnet + ) + + return AgentModelConfig(**config_data) + + except Exception as e: + logger.error(f"Error getting agent model config: {e}") + return None + + async def get_model_for_task( + self, + agent_id: str, + task_capabilities: List[ModelCapability], + cost_limit: Optional[float] = None, + ) -> Optional[str]: + """Select best model for a task based on capabilities and cost""" + try: + agent_config = await self.get_agent_model_config(agent_id) + if not agent_config: + return None + + # Check if primary model supports required capabilities + primary_model = self.available_models.get(agent_config.primary_model) + if primary_model and all( + cap in primary_model.capabilities for cap in task_capabilities + ): + # Check cost limit if specified + if ( + cost_limit is None + or self._estimate_task_cost(primary_model, 1000) <= cost_limit + ): + return agent_config.primary_model + + # Try fallback models + for model_id in agent_config.fallback_models: + model = self.available_models.get(model_id) + if model and all( + cap in model.capabilities for cap in task_capabilities + ): + if ( + cost_limit is None + or self._estimate_task_cost(model, 1000) <= cost_limit + ): + return model_id + + # No suitable model found + logger.warning( + f"No suitable model found for agent {agent_id} with capabilities {task_capabilities}" + ) + return None + + except Exception as e: + logger.error(f"Error selecting model for task: {e}") + return None + + def _estimate_task_cost(self, model: ModelSpec, estimated_tokens: int) -> float: + """Estimate cost for a task with given token count""" + # Simple estimation assuming 70% input, 30% output tokens + input_tokens = int(estimated_tokens * 0.7) + output_tokens = int(estimated_tokens * 0.3) + + input_cost = (input_tokens / 1000) * model.cost_per_input_token + output_cost = (output_tokens / 1000) * model.cost_per_output_token + + return input_cost + output_cost + + async def get_available_models( + self, + organization_id: str, + provider: Optional[ModelProvider] = None, + capabilities: Optional[List[ModelCapability]] = None, + ) -> List[Dict[str, Any]]: + """Get available models with provider credential validation""" + available = [] + + for model_id, model in self.available_models.items(): + # Filter by provider if specified + if provider and model.provider != provider: + continue + + # Filter by capabilities if specified + if capabilities and not all( + cap in model.capabilities for cap in capabilities + ): + continue + + # Check if organization has credentials for this provider + has_credentials = ( + await self.get_provider_credentials(organization_id, model.provider) + is not None + ) + + model_info = { + "model_id": model_id, + "provider": model.provider.value, + "name": model.name, + "description": model.description, + "capabilities": [cap.value for cap in model.capabilities], + "context_window": model.context_window, + "max_output_tokens": model.max_output_tokens, + "cost_per_input_token": model.cost_per_input_token, + "cost_per_output_token": model.cost_per_output_token, + "supports_streaming": model.supports_streaming, + "supports_function_calling": model.supports_function_calling, + "supports_vision": model.supports_vision, + "supports_json_mode": model.supports_json_mode, + "recommended_use_cases": model.recommended_use_cases, + "has_credentials": has_credentials, + "available": has_credentials, + } + + available.append(model_info) + + return available + + async def get_organization_model_usage( + self, organization_id: str, days: int = 30 + ) -> Dict[str, Any]: + """Get model usage statistics for an organization""" + try: + usage_data = await DatabaseManager.get_model_usage_stats( + organization_id, days + ) + + return { + "organization_id": organization_id, + "period_days": days, + "total_requests": usage_data.get("total_requests", 0), + "total_tokens": usage_data.get("total_tokens", 0), + "total_cost": usage_data.get("total_cost", 0.0), + "model_breakdown": usage_data.get("model_breakdown", {}), + "agent_breakdown": usage_data.get("agent_breakdown", {}), + "daily_usage": usage_data.get("daily_usage", []), + } + + except Exception as e: + logger.error(f"Error getting model usage: {e}") + return {} + + async def estimate_task_cost( + self, agent_id: str, task_description: str, estimated_complexity: str = "medium" + ) -> Dict[str, Any]: + """Estimate cost for a task execution""" + try: + agent_config = await self.get_agent_model_config(agent_id) + if not agent_config: + return {"error": "Agent configuration not found"} + + model = self.available_models.get(agent_config.primary_model) + if not model: + return {"error": "Model specification not found"} + + # Estimate token usage based on complexity + token_estimates = { + "low": 2000, + "medium": 5000, + "high": 10000, + "very_high": 20000, + } + + estimated_tokens = token_estimates.get(estimated_complexity, 5000) + estimated_cost = self._estimate_task_cost(model, estimated_tokens) + + return { + "agent_id": agent_id, + "model": model.model_id, + "estimated_tokens": estimated_tokens, + "estimated_cost_usd": round(estimated_cost, 4), + "complexity": estimated_complexity, + "cost_breakdown": { + "input_cost": (estimated_tokens * 0.7 / 1000) + * model.cost_per_input_token, + "output_cost": (estimated_tokens * 0.3 / 1000) + * model.cost_per_output_token, + }, + } + + except Exception as e: + logger.error(f"Error estimating task cost: {e}") + return {"error": str(e)} + + +# Global instance +model_config_manager = ModelConfigurationManager() diff --git a/services/orchestrator/multi_agent_coordinator.py b/services/orchestrator/multi_agent_coordinator.py index f1951ee..3c80db9 100644 --- a/services/orchestrator/multi_agent_coordinator.py +++ b/services/orchestrator/multi_agent_coordinator.py @@ -1,938 +1,938 @@ -""" -Multi-Agent Coordination System for FuzeAgent - -Enables multiple agents to collaborate on complex tasks through: -- Task decomposition and delegation -- Agent communication and synchronization -- Dependency management -- Result aggregation -- Conflict resolution -- Progress monitoring across agent teams - -This system allows for autonomous coordination of development teams -where agents can request help, delegate subtasks, and coordinate -work without human intervention. -""" - -import asyncio -import json -import logging -import uuid -from dataclasses import dataclass, field -from datetime import datetime, timedelta -from enum import Enum -from typing import Any, Dict, List, Optional, Set, Tuple - -from .database import DatabaseManager -from .task_execution_engine import TaskExecutionEngine, TaskStatus - -logger = logging.getLogger(__name__) - - -class CoordinationMode(str, Enum): - SEQUENTIAL = "sequential" # Tasks executed one after another - PARALLEL = "parallel" # Tasks executed simultaneously - HIERARCHICAL = "hierarchical" # Manager delegates to subordinates - COLLABORATIVE = "collaborative" # Agents work together on shared task - - -class AgentRole(str, Enum): - COORDINATOR = "coordinator" # Leads the coordination - PARTICIPANT = "participant" # Participates in coordination - OBSERVER = "observer" # Observes but doesn't execute - - -class CoordinationStatus(str, Enum): - INITIALIZING = "initializing" - PLANNING = "planning" - EXECUTING = "executing" - SYNCHRONIZING = "synchronizing" - REVIEWING = "reviewing" - COMPLETED = "completed" - FAILED = "failed" - CANCELLED = "cancelled" - - -@dataclass -class AgentCapability: - """Represents an agent's capability""" - - skill: str - proficiency: float # 0.0 to 1.0 - availability: bool - current_load: float # 0.0 to 1.0 - - -@dataclass -class TaskDependency: - """Represents a dependency between tasks""" - - dependent_task_id: str - prerequisite_task_id: str - dependency_type: str # "blocking", "soft", "informational" - - -@dataclass -class CoordinationPlan: - """Represents a plan for multi-agent coordination""" - - plan_id: str - root_task_id: str - coordination_mode: CoordinationMode - participating_agents: List[str] - task_assignments: Dict[str, str] # task_id -> agent_id - dependencies: List[TaskDependency] - estimated_completion: datetime - created_at: datetime - - -@dataclass -class AgentCommunication: - """Represents communication between agents""" - - communication_id: str - from_agent_id: str - to_agent_id: str - message_type: str # "request", "response", "notification", "question" - content: str - metadata: Dict[str, Any] - timestamp: datetime - response_id: Optional[str] = None - - -@dataclass -class CoordinationSession: - """Represents an active multi-agent coordination session""" - - session_id: str - root_task_id: str - coordinator_agent_id: str - participating_agents: Set[str] - coordination_mode: CoordinationMode - status: CoordinationStatus - plan: Optional[CoordinationPlan] - communications: List[AgentCommunication] = field(default_factory=list) - subtasks: Dict[str, str] = field(default_factory=dict) # subtask_id -> agent_id - started_at: datetime = field(default_factory=datetime.now) - completed_at: Optional[datetime] = None - result: Optional[Dict[str, Any]] = None - - -class MultiAgentCoordinator: - """ - Orchestrates multi-agent collaboration for complex tasks. - - Features: - - Automatic task decomposition - - Agent capability matching - - Dynamic load balancing - - Inter-agent communication - - Dependency resolution - - Progress synchronization - - Conflict resolution - """ - - def __init__(self, task_execution_engine: TaskExecutionEngine): - self.task_execution_engine = task_execution_engine - self.active_sessions: Dict[str, CoordinationSession] = {} - self.agent_capabilities: Dict[str, List[AgentCapability]] = {} - self.running = False - self.coordination_workers: List[asyncio.Task] = [] - - # Configuration - self.max_concurrent_coordinations = 10 - self.communication_timeout = 300 # 5 minutes - self.synchronization_interval = 30 # 30 seconds - - async def start(self): - """Start the multi-agent coordinator""" - logger.info("Starting MultiAgentCoordinator") - self.running = True - - # Start coordination workers - self.coordination_workers = [ - asyncio.create_task(self._coordination_worker()), - asyncio.create_task(self._communication_worker()), - asyncio.create_task(self._synchronization_worker()), - ] - - logger.info("MultiAgentCoordinator started") - - async def stop(self): - """Stop the multi-agent coordinator""" - logger.info("Stopping MultiAgentCoordinator") - self.running = False - - # Cancel workers - for worker in self.coordination_workers: - worker.cancel() - - try: - await asyncio.gather(*self.coordination_workers, return_exceptions=True) - except Exception as e: - logger.error(f"Error stopping coordination workers: {e}") - - # Clean up active sessions - for session_id in list(self.active_sessions.keys()): - await self._cleanup_session(session_id) - - logger.info("MultiAgentCoordinator stopped") - - async def initiate_coordination( - self, - task_id: str, - coordination_mode: CoordinationMode = CoordinationMode.COLLABORATIVE, - required_agents: Optional[List[str]] = None, - required_skills: Optional[List[str]] = None, - ) -> str: - """ - Initiate multi-agent coordination for a complex task. - - Args: - task_id: The root task to coordinate - coordination_mode: How agents should coordinate - required_agents: Specific agents to include - required_skills: Required skills for the task - - Returns: - Coordination session ID - """ - logger.info(f"Initiating coordination for task {task_id}") - - try: - # Get task information - task_data = await DatabaseManager.get_task(task_id) - if not task_data: - raise ValueError(f"Task {task_id} not found") - - # Analyze task complexity and determine if coordination is needed - complexity_analysis = await self._analyze_task_complexity(task_data) - - if not complexity_analysis["requires_coordination"]: - logger.info(f"Task {task_id} does not require coordination") - return None - - # Find suitable agents - if required_agents: - selected_agents = required_agents - else: - selected_agents = await self._select_agents_for_task( - task_data, required_skills, complexity_analysis - ) - - if len(selected_agents) < 2: - logger.warning( - f"Not enough agents available for coordination: {len(selected_agents)}" - ) - return None - - # Determine coordinator agent (first agent or most experienced) - coordinator_agent_id = await self._select_coordinator( - selected_agents, task_data - ) - - # Create coordination session - session_id = str(uuid.uuid4()) - session = CoordinationSession( - session_id=session_id, - root_task_id=task_id, - coordinator_agent_id=coordinator_agent_id, - participating_agents=set(selected_agents), - coordination_mode=coordination_mode, - status=CoordinationStatus.INITIALIZING, - ) - - self.active_sessions[session_id] = session - - logger.info( - f"Created coordination session {session_id} with {len(selected_agents)} agents" - ) - return session_id - - except Exception as e: - logger.error(f"Failed to initiate coordination for task {task_id}: {e}") - raise - - async def get_coordination_status( - self, session_id: str - ) -> Optional[Dict[str, Any]]: - """Get status of a coordination session""" - session = self.active_sessions.get(session_id) - if not session: - return None - - # Get subtask statuses - subtask_statuses = {} - for subtask_id, agent_id in session.subtasks.items(): - status = await self.task_execution_engine.get_execution_status(subtask_id) - subtask_statuses[subtask_id] = { - "agent_id": agent_id, - "status": status.get("status", "unknown") if status else "unknown", - } - - return { - "session_id": session_id, - "root_task_id": session.root_task_id, - "coordinator": session.coordinator_agent_id, - "participating_agents": list(session.participating_agents), - "coordination_mode": session.coordination_mode.value, - "status": session.status.value, - "subtasks": subtask_statuses, - "communications_count": len(session.communications), - "started_at": session.started_at.isoformat(), - "completed_at": ( - session.completed_at.isoformat() if session.completed_at else None - ), - "estimated_completion": ( - session.plan.estimated_completion.isoformat() if session.plan else None - ), - } - - async def send_agent_communication( - self, - from_agent_id: str, - to_agent_id: str, - message_type: str, - content: str, - metadata: Optional[Dict[str, Any]] = None, - ) -> str: - """Send communication between agents""" - - communication_id = str(uuid.uuid4()) - communication = AgentCommunication( - communication_id=communication_id, - from_agent_id=from_agent_id, - to_agent_id=to_agent_id, - message_type=message_type, - content=content, - metadata=metadata or {}, - timestamp=datetime.now(), - ) - - # Find coordination session for these agents - session = self._find_session_by_agents([from_agent_id, to_agent_id]) - if session: - session.communications.append(communication) - - # Store in database - await self._store_communication(communication) - - logger.info( - f"Agent communication {communication_id}: {from_agent_id} -> {to_agent_id}" - ) - return communication_id - - async def cancel_coordination(self, session_id: str) -> bool: - """Cancel an active coordination session""" - session = self.active_sessions.get(session_id) - if not session: - return False - - try: - # Cancel all subtasks - for subtask_id in session.subtasks.keys(): - await self.task_execution_engine.cancel_task_execution(subtask_id) - - # Update session status - session.status = CoordinationStatus.CANCELLED - session.completed_at = datetime.now() - - # Cleanup - await self._cleanup_session(session_id) - - logger.info(f"Cancelled coordination session {session_id}") - return True - - except Exception as e: - logger.error(f"Error cancelling coordination session {session_id}: {e}") - return False - - # Private methods - - async def _coordination_worker(self): - """Main coordination worker that manages session lifecycle""" - while self.running: - try: - # Process sessions that need attention - for session_id, session in list(self.active_sessions.items()): - try: - await self._process_coordination_session(session) - except Exception as e: - logger.error( - f"Error processing coordination session {session_id}: {e}" - ) - - await asyncio.sleep(5) # Check every 5 seconds - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Error in coordination worker: {e}") - await asyncio.sleep(10) - - async def _communication_worker(self): - """Worker that handles inter-agent communications""" - while self.running: - try: - # Process pending communications - for session in self.active_sessions.values(): - for comm in session.communications: - if not comm.response_id and comm.message_type == "request": - # Check if communication has timed out - if ( - datetime.now() - comm.timestamp - ).total_seconds() > self.communication_timeout: - await self._handle_communication_timeout(session, comm) - - await asyncio.sleep(10) # Check every 10 seconds - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Error in communication worker: {e}") - await asyncio.sleep(10) - - async def _synchronization_worker(self): - """Worker that synchronizes coordination sessions""" - while self.running: - try: - for session_id, session in list(self.active_sessions.items()): - if session.status == CoordinationStatus.EXECUTING: - await self._synchronize_session(session) - - await asyncio.sleep(self.synchronization_interval) - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Error in synchronization worker: {e}") - await asyncio.sleep(self.synchronization_interval) - - async def _process_coordination_session(self, session: CoordinationSession): - """Process a coordination session based on its current status""" - - if session.status == CoordinationStatus.INITIALIZING: - await self._initialize_session(session) - elif session.status == CoordinationStatus.PLANNING: - await self._plan_coordination(session) - elif session.status == CoordinationStatus.EXECUTING: - await self._monitor_execution(session) - elif session.status == CoordinationStatus.REVIEWING: - await self._review_coordination(session) - - async def _initialize_session(self, session: CoordinationSession): - """Initialize a coordination session""" - try: - # Get detailed task information - task_data = await DatabaseManager.get_task(session.root_task_id) - - # Update agent capabilities - await self._update_agent_capabilities(list(session.participating_agents)) - - # Move to planning phase - session.status = CoordinationStatus.PLANNING - - logger.info(f"Initialized coordination session {session.session_id}") - - except Exception as e: - logger.error(f"Error initializing session {session.session_id}: {e}") - session.status = CoordinationStatus.FAILED - - async def _plan_coordination(self, session: CoordinationSession): - """Create coordination plan""" - try: - # Get task data - task_data = await DatabaseManager.get_task(session.root_task_id) - - # Decompose task into subtasks - subtasks = await self._decompose_task(task_data, session.coordination_mode) - - # Assign agents to subtasks - assignments = await self._assign_agents_to_subtasks( - subtasks, list(session.participating_agents) - ) - - # Create dependencies - dependencies = await self._create_task_dependencies( - subtasks, session.coordination_mode - ) - - # Estimate completion time - estimated_completion = await self._estimate_coordination_completion( - subtasks, assignments, dependencies - ) - - # Create coordination plan - plan = CoordinationPlan( - plan_id=str(uuid.uuid4()), - root_task_id=session.root_task_id, - coordination_mode=session.coordination_mode, - participating_agents=list(session.participating_agents), - task_assignments=assignments, - dependencies=dependencies, - estimated_completion=estimated_completion, - created_at=datetime.now(), - ) - - session.plan = plan - session.status = CoordinationStatus.EXECUTING - - # Create subtasks in database and start execution - for subtask_data in subtasks: - subtask_id = await self._create_subtask(subtask_data, assignments) - session.subtasks[subtask_id] = assignments[subtask_data["id"]] - - # Start subtask execution - await self.task_execution_engine.start_task_execution(subtask_id) - - logger.info( - f"Created coordination plan for session {session.session_id} with {len(subtasks)} subtasks" - ) - - except Exception as e: - logger.error(f"Error planning coordination {session.session_id}: {e}") - session.status = CoordinationStatus.FAILED - - async def _monitor_execution(self, session: CoordinationSession): - """Monitor execution of coordinated tasks""" - try: - # Check status of all subtasks - completed_subtasks = 0 - failed_subtasks = 0 - - for subtask_id in session.subtasks.keys(): - status = await self.task_execution_engine.get_execution_status( - subtask_id - ) - if status: - if status.get("status") == "completed": - completed_subtasks += 1 - elif status.get("status") == "failed": - failed_subtasks += 1 - - total_subtasks = len(session.subtasks) - - # Check if coordination is complete - if completed_subtasks == total_subtasks: - session.status = CoordinationStatus.REVIEWING - elif failed_subtasks > 0: - # Handle failures - await self._handle_coordination_failures(session) - - except Exception as e: - logger.error(f"Error monitoring execution {session.session_id}: {e}") - - async def _review_coordination(self, session: CoordinationSession): - """Review completed coordination and aggregate results""" - try: - # Collect results from all subtasks - results = {} - for subtask_id, agent_id in session.subtasks.items(): - status = await self.task_execution_engine.get_execution_status( - subtask_id - ) - if status: - results[subtask_id] = { - "agent_id": agent_id, - "status": status.get("status"), - "result": status.get("result"), - } - - # Aggregate results - coordination_result = await self._aggregate_coordination_results(results) - - # Complete coordination - session.status = CoordinationStatus.COMPLETED - session.completed_at = datetime.now() - session.result = coordination_result - - # Update root task status - await DatabaseManager.update_task_status( - session.root_task_id, "completed", coordination_result - ) - - logger.info(f"Completed coordination session {session.session_id}") - - # Schedule cleanup - asyncio.create_task(self._cleanup_session(session.session_id)) - - except Exception as e: - logger.error(f"Error reviewing coordination {session.session_id}: {e}") - session.status = CoordinationStatus.FAILED - - async def _analyze_task_complexity( - self, task_data: Dict[str, Any] - ) -> Dict[str, Any]: - """Analyze task complexity to determine if coordination is needed""" - - description = task_data.get("description", "") - title = task_data.get("title", "") - - # Simple heuristics for complexity analysis - complexity_indicators = [ - "multiple components" in description.lower(), - "frontend and backend" in description.lower(), - "database and api" in description.lower(), - "testing and deployment" in description.lower(), - len(description.split()) > 50, # Long description - "integrate" in description.lower(), - "coordinate" in description.lower(), - "collaborate" in description.lower(), - ] - - complexity_score = sum(complexity_indicators) / len(complexity_indicators) - - return { - "requires_coordination": complexity_score > 0.3, - "complexity_score": complexity_score, - "estimated_agents_needed": min(max(2, int(complexity_score * 5)), 5), - "estimated_duration_hours": max(4, int(complexity_score * 24)), - } - - async def _select_agents_for_task( - self, - task_data: Dict[str, Any], - required_skills: Optional[List[str]], - complexity_analysis: Dict[str, Any], - ) -> List[str]: - """Select appropriate agents for the task""" - - # Get all available agents - all_agents = await DatabaseManager.get_all_agents() - - # Filter by availability and skills - suitable_agents = [] - for agent in all_agents: - if agent["status"] == "available": - agent_skills = agent.get("config", {}).get("tools", []) - - # Check skill match - if required_skills: - skill_match = any( - skill in agent_skills for skill in required_skills - ) - else: - skill_match = True - - if skill_match: - suitable_agents.append(agent["id"]) - - # Select optimal number of agents - max_agents = complexity_analysis.get("estimated_agents_needed", 3) - return suitable_agents[:max_agents] - - async def _select_coordinator( - self, agents: List[str], task_data: Dict[str, Any] - ) -> str: - """Select the coordinator agent from available agents""" - - # For now, select the first agent as coordinator - # In production, this would consider agent experience, current load, etc. - return agents[0] - - async def _decompose_task( - self, task_data: Dict[str, Any], mode: CoordinationMode - ) -> List[Dict[str, Any]]: - """Decompose a complex task into subtasks""" - - # Simple task decomposition based on common patterns - description = task_data.get("description", "") - title = task_data.get("title", "") - - subtasks = [] - - # Common subtask patterns - if "frontend" in description.lower() or "ui" in description.lower(): - subtasks.append( - { - "id": f"frontend-{uuid.uuid4()}", - "title": f"Frontend Implementation - {title}", - "description": f"Implement frontend components for: {description}", - "type": "frontend_development", - "estimated_hours": 4, - } - ) - - if "backend" in description.lower() or "api" in description.lower(): - subtasks.append( - { - "id": f"backend-{uuid.uuid4()}", - "title": f"Backend Implementation - {title}", - "description": f"Implement backend services for: {description}", - "type": "backend_development", - "estimated_hours": 6, - } - ) - - if "database" in description.lower() or "data" in description.lower(): - subtasks.append( - { - "id": f"database-{uuid.uuid4()}", - "title": f"Database Design - {title}", - "description": f"Design and implement database schema for: {description}", - "type": "database_development", - "estimated_hours": 3, - } - ) - - if "test" in description.lower(): - subtasks.append( - { - "id": f"testing-{uuid.uuid4()}", - "title": f"Testing - {title}", - "description": f"Create comprehensive tests for: {description}", - "type": "testing", - "estimated_hours": 4, - } - ) - - # If no specific subtasks identified, create generic subtasks - if not subtasks: - subtasks = [ - { - "id": f"analysis-{uuid.uuid4()}", - "title": f"Analysis - {title}", - "description": f"Analyze requirements for: {description}", - "type": "analysis", - "estimated_hours": 2, - }, - { - "id": f"implementation-{uuid.uuid4()}", - "title": f"Implementation - {title}", - "description": f"Implement solution for: {description}", - "type": "implementation", - "estimated_hours": 6, - }, - { - "id": f"review-{uuid.uuid4()}", - "title": f"Review - {title}", - "description": f"Review and validate solution for: {description}", - "type": "review", - "estimated_hours": 2, - }, - ] - - return subtasks - - async def _assign_agents_to_subtasks( - self, subtasks: List[Dict[str, Any]], agents: List[str] - ) -> Dict[str, str]: - """Assign agents to subtasks based on capabilities""" - - assignments = {} - - # Get agent capabilities - agent_data = {} - for agent_id in agents: - agent = await DatabaseManager.get_agent(agent_id) - if agent: - agent_data[agent_id] = agent - - # Simple assignment based on agent type - for subtask in subtasks: - subtask_type = subtask.get("type", "") - best_agent = None - - # Match agent type to subtask type - for agent_id, agent in agent_data.items(): - agent_type = agent.get("type", "") - - if subtask_type.startswith("frontend") and "frontend" in agent_type: - best_agent = agent_id - break - elif subtask_type.startswith("backend") and "backend" in agent_type: - best_agent = agent_id - break - elif subtask_type.startswith("database") and "backend" in agent_type: - best_agent = agent_id - break - elif subtask_type == "testing" and "qa" in agent_type: - best_agent = agent_id - break - - # Fallback to first available agent - if not best_agent: - best_agent = agents[0] - - assignments[subtask["id"]] = best_agent - - return assignments - - async def _create_task_dependencies( - self, subtasks: List[Dict[str, Any]], mode: CoordinationMode - ) -> List[TaskDependency]: - """Create dependencies between subtasks""" - - dependencies = [] - - if mode == CoordinationMode.SEQUENTIAL: - # Create sequential dependencies - for i in range(1, len(subtasks)): - dependencies.append( - TaskDependency( - dependent_task_id=subtasks[i]["id"], - prerequisite_task_id=subtasks[i - 1]["id"], - dependency_type="blocking", - ) - ) - - elif mode == CoordinationMode.HIERARCHICAL: - # Analysis task should complete before implementation - analysis_tasks = [t for t in subtasks if "analysis" in t["type"]] - implementation_tasks = [ - t for t in subtasks if "implementation" in t["type"] - ] - - for impl_task in implementation_tasks: - for analysis_task in analysis_tasks: - dependencies.append( - TaskDependency( - dependent_task_id=impl_task["id"], - prerequisite_task_id=analysis_task["id"], - dependency_type="blocking", - ) - ) - - # PARALLEL and COLLABORATIVE modes have no strict dependencies - - return dependencies - - async def _estimate_coordination_completion( - self, - subtasks: List[Dict[str, Any]], - assignments: Dict[str, str], - dependencies: List[TaskDependency], - ) -> datetime: - """Estimate when coordination will complete""" - - if not dependencies: - # Parallel execution - completion time is max of all subtasks - max_hours = max(subtask.get("estimated_hours", 4) for subtask in subtasks) - else: - # Sequential/dependent execution - sum of critical path - total_hours = sum(subtask.get("estimated_hours", 4) for subtask in subtasks) - max_hours = min(total_hours, 24) # Cap at 24 hours - - return datetime.now() + timedelta(hours=max_hours) - - async def _create_subtask( - self, subtask_data: Dict[str, Any], assignments: Dict[str, str] - ) -> str: - """Create a subtask in the database""" - - subtask_id = str(uuid.uuid4()) - agent_id = assignments.get(subtask_data["id"]) - - # Create task in database - await DatabaseManager.create_task( - task_id=subtask_id, - title=subtask_data["title"], - description=subtask_data["description"], - assigned_to=agent_id, - priority="medium", - metadata={ - "coordination_subtask": True, - "parent_task_type": subtask_data.get("type"), - "estimated_hours": subtask_data.get("estimated_hours", 4), - }, - ) - - return subtask_id - - def _find_session_by_agents( - self, agent_ids: List[str] - ) -> Optional[CoordinationSession]: - """Find coordination session that includes the specified agents""" - for session in self.active_sessions.values(): - if any(agent_id in session.participating_agents for agent_id in agent_ids): - return session - return None - - async def _store_communication(self, communication: AgentCommunication): - """Store agent communication in database""" - # This would store the communication in the database - # For now, just log it - logger.info( - f"Agent communication: {communication.from_agent_id} -> {communication.to_agent_id}: {communication.content}" - ) - - async def _handle_communication_timeout( - self, session: CoordinationSession, communication: AgentCommunication - ): - """Handle communication timeout""" - logger.warning( - f"Communication timeout in session {session.session_id}: {communication.communication_id}" - ) - - # Could implement retry logic or escalation here - - async def _synchronize_session(self, session: CoordinationSession): - """Synchronize coordination session state""" - # Check if any agents need help or coordination - # Update session status based on subtask progress - # Handle any conflicts or issues - pass - - async def _handle_coordination_failures(self, session: CoordinationSession): - """Handle failures in coordination""" - logger.warning( - f"Handling failures in coordination session {session.session_id}" - ) - - # Could implement retry logic, reassignment, or escalation - session.status = CoordinationStatus.FAILED - - async def _aggregate_coordination_results( - self, results: Dict[str, Dict[str, Any]] - ) -> Dict[str, Any]: - """Aggregate results from all coordinated subtasks""" - - successful_subtasks = [ - r for r in results.values() if r["status"] == "completed" - ] - - return { - "coordination_completed": True, - "total_subtasks": len(results), - "successful_subtasks": len(successful_subtasks), - "failed_subtasks": len(results) - len(successful_subtasks), - "results": results, - "completion_time": datetime.now().isoformat(), - } - - async def _update_agent_capabilities(self, agent_ids: List[str]): - """Update cached agent capabilities""" - for agent_id in agent_ids: - agent_data = await DatabaseManager.get_agent(agent_id) - if agent_data: - # Extract capabilities from agent configuration - tools = agent_data.get("config", {}).get("tools", []) - capabilities = [ - AgentCapability( - skill=tool, - proficiency=0.8, # Default proficiency - availability=agent_data.get("status") == "available", - current_load=0.5, # Default load - ) - for tool in tools - ] - self.agent_capabilities[agent_id] = capabilities - - async def _cleanup_session(self, session_id: str): - """Clean up completed coordination session""" - session = self.active_sessions.pop(session_id, None) - if session: - logger.info(f"Cleaned up coordination session {session_id}") - - -# Integration with existing TaskExecutionEngine -def integrate_multi_agent_coordination( - task_execution_engine: TaskExecutionEngine, -) -> MultiAgentCoordinator: - """Create and integrate multi-agent coordinator with task execution engine""" - coordinator = MultiAgentCoordinator(task_execution_engine) - - # Add coordination capabilities to task execution engine - task_execution_engine.multi_agent_coordinator = coordinator - - return coordinator +""" +Multi-Agent Coordination System for FuzeAgent + +Enables multiple agents to collaborate on complex tasks through: +- Task decomposition and delegation +- Agent communication and synchronization +- Dependency management +- Result aggregation +- Conflict resolution +- Progress monitoring across agent teams + +This system allows for autonomous coordination of development teams +where agents can request help, delegate subtasks, and coordinate +work without human intervention. +""" + +import asyncio +import json +import logging +import uuid +from dataclasses import dataclass, field +from datetime import datetime, timedelta +from enum import Enum +from typing import Any, Dict, List, Optional, Set, Tuple + +from .database import DatabaseManager +from .task_execution_engine import TaskExecutionEngine, TaskStatus + +logger = logging.getLogger(__name__) + + +class CoordinationMode(str, Enum): + SEQUENTIAL = "sequential" # Tasks executed one after another + PARALLEL = "parallel" # Tasks executed simultaneously + HIERARCHICAL = "hierarchical" # Manager delegates to subordinates + COLLABORATIVE = "collaborative" # Agents work together on shared task + + +class AgentRole(str, Enum): + COORDINATOR = "coordinator" # Leads the coordination + PARTICIPANT = "participant" # Participates in coordination + OBSERVER = "observer" # Observes but doesn't execute + + +class CoordinationStatus(str, Enum): + INITIALIZING = "initializing" + PLANNING = "planning" + EXECUTING = "executing" + SYNCHRONIZING = "synchronizing" + REVIEWING = "reviewing" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +@dataclass +class AgentCapability: + """Represents an agent's capability""" + + skill: str + proficiency: float # 0.0 to 1.0 + availability: bool + current_load: float # 0.0 to 1.0 + + +@dataclass +class TaskDependency: + """Represents a dependency between tasks""" + + dependent_task_id: str + prerequisite_task_id: str + dependency_type: str # "blocking", "soft", "informational" + + +@dataclass +class CoordinationPlan: + """Represents a plan for multi-agent coordination""" + + plan_id: str + root_task_id: str + coordination_mode: CoordinationMode + participating_agents: List[str] + task_assignments: Dict[str, str] # task_id -> agent_id + dependencies: List[TaskDependency] + estimated_completion: datetime + created_at: datetime + + +@dataclass +class AgentCommunication: + """Represents communication between agents""" + + communication_id: str + from_agent_id: str + to_agent_id: str + message_type: str # "request", "response", "notification", "question" + content: str + metadata: Dict[str, Any] + timestamp: datetime + response_id: Optional[str] = None + + +@dataclass +class CoordinationSession: + """Represents an active multi-agent coordination session""" + + session_id: str + root_task_id: str + coordinator_agent_id: str + participating_agents: Set[str] + coordination_mode: CoordinationMode + status: CoordinationStatus + plan: Optional[CoordinationPlan] + communications: List[AgentCommunication] = field(default_factory=list) + subtasks: Dict[str, str] = field(default_factory=dict) # subtask_id -> agent_id + started_at: datetime = field(default_factory=datetime.now) + completed_at: Optional[datetime] = None + result: Optional[Dict[str, Any]] = None + + +class MultiAgentCoordinator: + """ + Orchestrates multi-agent collaboration for complex tasks. + + Features: + - Automatic task decomposition + - Agent capability matching + - Dynamic load balancing + - Inter-agent communication + - Dependency resolution + - Progress synchronization + - Conflict resolution + """ + + def __init__(self, task_execution_engine: TaskExecutionEngine): + self.task_execution_engine = task_execution_engine + self.active_sessions: Dict[str, CoordinationSession] = {} + self.agent_capabilities: Dict[str, List[AgentCapability]] = {} + self.running = False + self.coordination_workers: List[asyncio.Task] = [] + + # Configuration + self.max_concurrent_coordinations = 10 + self.communication_timeout = 300 # 5 minutes + self.synchronization_interval = 30 # 30 seconds + + async def start(self): + """Start the multi-agent coordinator""" + logger.info("Starting MultiAgentCoordinator") + self.running = True + + # Start coordination workers + self.coordination_workers = [ + asyncio.create_task(self._coordination_worker()), + asyncio.create_task(self._communication_worker()), + asyncio.create_task(self._synchronization_worker()), + ] + + logger.info("MultiAgentCoordinator started") + + async def stop(self): + """Stop the multi-agent coordinator""" + logger.info("Stopping MultiAgentCoordinator") + self.running = False + + # Cancel workers + for worker in self.coordination_workers: + worker.cancel() + + try: + await asyncio.gather(*self.coordination_workers, return_exceptions=True) + except Exception as e: + logger.error(f"Error stopping coordination workers: {e}") + + # Clean up active sessions + for session_id in list(self.active_sessions.keys()): + await self._cleanup_session(session_id) + + logger.info("MultiAgentCoordinator stopped") + + async def initiate_coordination( + self, + task_id: str, + coordination_mode: CoordinationMode = CoordinationMode.COLLABORATIVE, + required_agents: Optional[List[str]] = None, + required_skills: Optional[List[str]] = None, + ) -> str: + """ + Initiate multi-agent coordination for a complex task. + + Args: + task_id: The root task to coordinate + coordination_mode: How agents should coordinate + required_agents: Specific agents to include + required_skills: Required skills for the task + + Returns: + Coordination session ID + """ + logger.info(f"Initiating coordination for task {task_id}") + + try: + # Get task information + task_data = await DatabaseManager.get_task(task_id) + if not task_data: + raise ValueError(f"Task {task_id} not found") + + # Analyze task complexity and determine if coordination is needed + complexity_analysis = await self._analyze_task_complexity(task_data) + + if not complexity_analysis["requires_coordination"]: + logger.info(f"Task {task_id} does not require coordination") + return None + + # Find suitable agents + if required_agents: + selected_agents = required_agents + else: + selected_agents = await self._select_agents_for_task( + task_data, required_skills, complexity_analysis + ) + + if len(selected_agents) < 2: + logger.warning( + f"Not enough agents available for coordination: {len(selected_agents)}" + ) + return None + + # Determine coordinator agent (first agent or most experienced) + coordinator_agent_id = await self._select_coordinator( + selected_agents, task_data + ) + + # Create coordination session + session_id = str(uuid.uuid4()) + session = CoordinationSession( + session_id=session_id, + root_task_id=task_id, + coordinator_agent_id=coordinator_agent_id, + participating_agents=set(selected_agents), + coordination_mode=coordination_mode, + status=CoordinationStatus.INITIALIZING, + ) + + self.active_sessions[session_id] = session + + logger.info( + f"Created coordination session {session_id} with {len(selected_agents)} agents" + ) + return session_id + + except Exception as e: + logger.error(f"Failed to initiate coordination for task {task_id}: {e}") + raise + + async def get_coordination_status( + self, session_id: str + ) -> Optional[Dict[str, Any]]: + """Get status of a coordination session""" + session = self.active_sessions.get(session_id) + if not session: + return None + + # Get subtask statuses + subtask_statuses = {} + for subtask_id, agent_id in session.subtasks.items(): + status = await self.task_execution_engine.get_execution_status(subtask_id) + subtask_statuses[subtask_id] = { + "agent_id": agent_id, + "status": status.get("status", "unknown") if status else "unknown", + } + + return { + "session_id": session_id, + "root_task_id": session.root_task_id, + "coordinator": session.coordinator_agent_id, + "participating_agents": list(session.participating_agents), + "coordination_mode": session.coordination_mode.value, + "status": session.status.value, + "subtasks": subtask_statuses, + "communications_count": len(session.communications), + "started_at": session.started_at.isoformat(), + "completed_at": ( + session.completed_at.isoformat() if session.completed_at else None + ), + "estimated_completion": ( + session.plan.estimated_completion.isoformat() if session.plan else None + ), + } + + async def send_agent_communication( + self, + from_agent_id: str, + to_agent_id: str, + message_type: str, + content: str, + metadata: Optional[Dict[str, Any]] = None, + ) -> str: + """Send communication between agents""" + + communication_id = str(uuid.uuid4()) + communication = AgentCommunication( + communication_id=communication_id, + from_agent_id=from_agent_id, + to_agent_id=to_agent_id, + message_type=message_type, + content=content, + metadata=metadata or {}, + timestamp=datetime.now(), + ) + + # Find coordination session for these agents + session = self._find_session_by_agents([from_agent_id, to_agent_id]) + if session: + session.communications.append(communication) + + # Store in database + await self._store_communication(communication) + + logger.info( + f"Agent communication {communication_id}: {from_agent_id} -> {to_agent_id}" + ) + return communication_id + + async def cancel_coordination(self, session_id: str) -> bool: + """Cancel an active coordination session""" + session = self.active_sessions.get(session_id) + if not session: + return False + + try: + # Cancel all subtasks + for subtask_id in session.subtasks.keys(): + await self.task_execution_engine.cancel_task_execution(subtask_id) + + # Update session status + session.status = CoordinationStatus.CANCELLED + session.completed_at = datetime.now() + + # Cleanup + await self._cleanup_session(session_id) + + logger.info(f"Cancelled coordination session {session_id}") + return True + + except Exception as e: + logger.error(f"Error cancelling coordination session {session_id}: {e}") + return False + + # Private methods + + async def _coordination_worker(self): + """Main coordination worker that manages session lifecycle""" + while self.running: + try: + # Process sessions that need attention + for session_id, session in list(self.active_sessions.items()): + try: + await self._process_coordination_session(session) + except Exception as e: + logger.error( + f"Error processing coordination session {session_id}: {e}" + ) + + await asyncio.sleep(5) # Check every 5 seconds + + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"Error in coordination worker: {e}") + await asyncio.sleep(10) + + async def _communication_worker(self): + """Worker that handles inter-agent communications""" + while self.running: + try: + # Process pending communications + for session in self.active_sessions.values(): + for comm in session.communications: + if not comm.response_id and comm.message_type == "request": + # Check if communication has timed out + if ( + datetime.now() - comm.timestamp + ).total_seconds() > self.communication_timeout: + await self._handle_communication_timeout(session, comm) + + await asyncio.sleep(10) # Check every 10 seconds + + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"Error in communication worker: {e}") + await asyncio.sleep(10) + + async def _synchronization_worker(self): + """Worker that synchronizes coordination sessions""" + while self.running: + try: + for session_id, session in list(self.active_sessions.items()): + if session.status == CoordinationStatus.EXECUTING: + await self._synchronize_session(session) + + await asyncio.sleep(self.synchronization_interval) + + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"Error in synchronization worker: {e}") + await asyncio.sleep(self.synchronization_interval) + + async def _process_coordination_session(self, session: CoordinationSession): + """Process a coordination session based on its current status""" + + if session.status == CoordinationStatus.INITIALIZING: + await self._initialize_session(session) + elif session.status == CoordinationStatus.PLANNING: + await self._plan_coordination(session) + elif session.status == CoordinationStatus.EXECUTING: + await self._monitor_execution(session) + elif session.status == CoordinationStatus.REVIEWING: + await self._review_coordination(session) + + async def _initialize_session(self, session: CoordinationSession): + """Initialize a coordination session""" + try: + # Get detailed task information + task_data = await DatabaseManager.get_task(session.root_task_id) + + # Update agent capabilities + await self._update_agent_capabilities(list(session.participating_agents)) + + # Move to planning phase + session.status = CoordinationStatus.PLANNING + + logger.info(f"Initialized coordination session {session.session_id}") + + except Exception as e: + logger.error(f"Error initializing session {session.session_id}: {e}") + session.status = CoordinationStatus.FAILED + + async def _plan_coordination(self, session: CoordinationSession): + """Create coordination plan""" + try: + # Get task data + task_data = await DatabaseManager.get_task(session.root_task_id) + + # Decompose task into subtasks + subtasks = await self._decompose_task(task_data, session.coordination_mode) + + # Assign agents to subtasks + assignments = await self._assign_agents_to_subtasks( + subtasks, list(session.participating_agents) + ) + + # Create dependencies + dependencies = await self._create_task_dependencies( + subtasks, session.coordination_mode + ) + + # Estimate completion time + estimated_completion = await self._estimate_coordination_completion( + subtasks, assignments, dependencies + ) + + # Create coordination plan + plan = CoordinationPlan( + plan_id=str(uuid.uuid4()), + root_task_id=session.root_task_id, + coordination_mode=session.coordination_mode, + participating_agents=list(session.participating_agents), + task_assignments=assignments, + dependencies=dependencies, + estimated_completion=estimated_completion, + created_at=datetime.now(), + ) + + session.plan = plan + session.status = CoordinationStatus.EXECUTING + + # Create subtasks in database and start execution + for subtask_data in subtasks: + subtask_id = await self._create_subtask(subtask_data, assignments) + session.subtasks[subtask_id] = assignments[subtask_data["id"]] + + # Start subtask execution + await self.task_execution_engine.start_task_execution(subtask_id) + + logger.info( + f"Created coordination plan for session {session.session_id} with {len(subtasks)} subtasks" + ) + + except Exception as e: + logger.error(f"Error planning coordination {session.session_id}: {e}") + session.status = CoordinationStatus.FAILED + + async def _monitor_execution(self, session: CoordinationSession): + """Monitor execution of coordinated tasks""" + try: + # Check status of all subtasks + completed_subtasks = 0 + failed_subtasks = 0 + + for subtask_id in session.subtasks.keys(): + status = await self.task_execution_engine.get_execution_status( + subtask_id + ) + if status: + if status.get("status") == "completed": + completed_subtasks += 1 + elif status.get("status") == "failed": + failed_subtasks += 1 + + total_subtasks = len(session.subtasks) + + # Check if coordination is complete + if completed_subtasks == total_subtasks: + session.status = CoordinationStatus.REVIEWING + elif failed_subtasks > 0: + # Handle failures + await self._handle_coordination_failures(session) + + except Exception as e: + logger.error(f"Error monitoring execution {session.session_id}: {e}") + + async def _review_coordination(self, session: CoordinationSession): + """Review completed coordination and aggregate results""" + try: + # Collect results from all subtasks + results = {} + for subtask_id, agent_id in session.subtasks.items(): + status = await self.task_execution_engine.get_execution_status( + subtask_id + ) + if status: + results[subtask_id] = { + "agent_id": agent_id, + "status": status.get("status"), + "result": status.get("result"), + } + + # Aggregate results + coordination_result = await self._aggregate_coordination_results(results) + + # Complete coordination + session.status = CoordinationStatus.COMPLETED + session.completed_at = datetime.now() + session.result = coordination_result + + # Update root task status + await DatabaseManager.update_task_status( + session.root_task_id, "completed", coordination_result + ) + + logger.info(f"Completed coordination session {session.session_id}") + + # Schedule cleanup + asyncio.create_task(self._cleanup_session(session.session_id)) + + except Exception as e: + logger.error(f"Error reviewing coordination {session.session_id}: {e}") + session.status = CoordinationStatus.FAILED + + async def _analyze_task_complexity( + self, task_data: Dict[str, Any] + ) -> Dict[str, Any]: + """Analyze task complexity to determine if coordination is needed""" + + description = task_data.get("description", "") + title = task_data.get("title", "") + + # Simple heuristics for complexity analysis + complexity_indicators = [ + "multiple components" in description.lower(), + "frontend and backend" in description.lower(), + "database and api" in description.lower(), + "testing and deployment" in description.lower(), + len(description.split()) > 50, # Long description + "integrate" in description.lower(), + "coordinate" in description.lower(), + "collaborate" in description.lower(), + ] + + complexity_score = sum(complexity_indicators) / len(complexity_indicators) + + return { + "requires_coordination": complexity_score > 0.3, + "complexity_score": complexity_score, + "estimated_agents_needed": min(max(2, int(complexity_score * 5)), 5), + "estimated_duration_hours": max(4, int(complexity_score * 24)), + } + + async def _select_agents_for_task( + self, + task_data: Dict[str, Any], + required_skills: Optional[List[str]], + complexity_analysis: Dict[str, Any], + ) -> List[str]: + """Select appropriate agents for the task""" + + # Get all available agents + all_agents = await DatabaseManager.get_all_agents() + + # Filter by availability and skills + suitable_agents = [] + for agent in all_agents: + if agent["status"] == "available": + agent_skills = agent.get("config", {}).get("tools", []) + + # Check skill match + if required_skills: + skill_match = any( + skill in agent_skills for skill in required_skills + ) + else: + skill_match = True + + if skill_match: + suitable_agents.append(agent["id"]) + + # Select optimal number of agents + max_agents = complexity_analysis.get("estimated_agents_needed", 3) + return suitable_agents[:max_agents] + + async def _select_coordinator( + self, agents: List[str], task_data: Dict[str, Any] + ) -> str: + """Select the coordinator agent from available agents""" + + # For now, select the first agent as coordinator + # In production, this would consider agent experience, current load, etc. + return agents[0] + + async def _decompose_task( + self, task_data: Dict[str, Any], mode: CoordinationMode + ) -> List[Dict[str, Any]]: + """Decompose a complex task into subtasks""" + + # Simple task decomposition based on common patterns + description = task_data.get("description", "") + title = task_data.get("title", "") + + subtasks = [] + + # Common subtask patterns + if "frontend" in description.lower() or "ui" in description.lower(): + subtasks.append( + { + "id": f"frontend-{uuid.uuid4()}", + "title": f"Frontend Implementation - {title}", + "description": f"Implement frontend components for: {description}", + "type": "frontend_development", + "estimated_hours": 4, + } + ) + + if "backend" in description.lower() or "api" in description.lower(): + subtasks.append( + { + "id": f"backend-{uuid.uuid4()}", + "title": f"Backend Implementation - {title}", + "description": f"Implement backend services for: {description}", + "type": "backend_development", + "estimated_hours": 6, + } + ) + + if "database" in description.lower() or "data" in description.lower(): + subtasks.append( + { + "id": f"database-{uuid.uuid4()}", + "title": f"Database Design - {title}", + "description": f"Design and implement database schema for: {description}", + "type": "database_development", + "estimated_hours": 3, + } + ) + + if "test" in description.lower(): + subtasks.append( + { + "id": f"testing-{uuid.uuid4()}", + "title": f"Testing - {title}", + "description": f"Create comprehensive tests for: {description}", + "type": "testing", + "estimated_hours": 4, + } + ) + + # If no specific subtasks identified, create generic subtasks + if not subtasks: + subtasks = [ + { + "id": f"analysis-{uuid.uuid4()}", + "title": f"Analysis - {title}", + "description": f"Analyze requirements for: {description}", + "type": "analysis", + "estimated_hours": 2, + }, + { + "id": f"implementation-{uuid.uuid4()}", + "title": f"Implementation - {title}", + "description": f"Implement solution for: {description}", + "type": "implementation", + "estimated_hours": 6, + }, + { + "id": f"review-{uuid.uuid4()}", + "title": f"Review - {title}", + "description": f"Review and validate solution for: {description}", + "type": "review", + "estimated_hours": 2, + }, + ] + + return subtasks + + async def _assign_agents_to_subtasks( + self, subtasks: List[Dict[str, Any]], agents: List[str] + ) -> Dict[str, str]: + """Assign agents to subtasks based on capabilities""" + + assignments = {} + + # Get agent capabilities + agent_data = {} + for agent_id in agents: + agent = await DatabaseManager.get_agent(agent_id) + if agent: + agent_data[agent_id] = agent + + # Simple assignment based on agent type + for subtask in subtasks: + subtask_type = subtask.get("type", "") + best_agent = None + + # Match agent type to subtask type + for agent_id, agent in agent_data.items(): + agent_type = agent.get("type", "") + + if subtask_type.startswith("frontend") and "frontend" in agent_type: + best_agent = agent_id + break + elif subtask_type.startswith("backend") and "backend" in agent_type: + best_agent = agent_id + break + elif subtask_type.startswith("database") and "backend" in agent_type: + best_agent = agent_id + break + elif subtask_type == "testing" and "qa" in agent_type: + best_agent = agent_id + break + + # Fallback to first available agent + if not best_agent: + best_agent = agents[0] + + assignments[subtask["id"]] = best_agent + + return assignments + + async def _create_task_dependencies( + self, subtasks: List[Dict[str, Any]], mode: CoordinationMode + ) -> List[TaskDependency]: + """Create dependencies between subtasks""" + + dependencies = [] + + if mode == CoordinationMode.SEQUENTIAL: + # Create sequential dependencies + for i in range(1, len(subtasks)): + dependencies.append( + TaskDependency( + dependent_task_id=subtasks[i]["id"], + prerequisite_task_id=subtasks[i - 1]["id"], + dependency_type="blocking", + ) + ) + + elif mode == CoordinationMode.HIERARCHICAL: + # Analysis task should complete before implementation + analysis_tasks = [t for t in subtasks if "analysis" in t["type"]] + implementation_tasks = [ + t for t in subtasks if "implementation" in t["type"] + ] + + for impl_task in implementation_tasks: + for analysis_task in analysis_tasks: + dependencies.append( + TaskDependency( + dependent_task_id=impl_task["id"], + prerequisite_task_id=analysis_task["id"], + dependency_type="blocking", + ) + ) + + # PARALLEL and COLLABORATIVE modes have no strict dependencies + + return dependencies + + async def _estimate_coordination_completion( + self, + subtasks: List[Dict[str, Any]], + assignments: Dict[str, str], + dependencies: List[TaskDependency], + ) -> datetime: + """Estimate when coordination will complete""" + + if not dependencies: + # Parallel execution - completion time is max of all subtasks + max_hours = max(subtask.get("estimated_hours", 4) for subtask in subtasks) + else: + # Sequential/dependent execution - sum of critical path + total_hours = sum(subtask.get("estimated_hours", 4) for subtask in subtasks) + max_hours = min(total_hours, 24) # Cap at 24 hours + + return datetime.now() + timedelta(hours=max_hours) + + async def _create_subtask( + self, subtask_data: Dict[str, Any], assignments: Dict[str, str] + ) -> str: + """Create a subtask in the database""" + + subtask_id = str(uuid.uuid4()) + agent_id = assignments.get(subtask_data["id"]) + + # Create task in database + await DatabaseManager.create_task( + task_id=subtask_id, + title=subtask_data["title"], + description=subtask_data["description"], + assigned_to=agent_id, + priority="medium", + metadata={ + "coordination_subtask": True, + "parent_task_type": subtask_data.get("type"), + "estimated_hours": subtask_data.get("estimated_hours", 4), + }, + ) + + return subtask_id + + def _find_session_by_agents( + self, agent_ids: List[str] + ) -> Optional[CoordinationSession]: + """Find coordination session that includes the specified agents""" + for session in self.active_sessions.values(): + if any(agent_id in session.participating_agents for agent_id in agent_ids): + return session + return None + + async def _store_communication(self, communication: AgentCommunication): + """Store agent communication in database""" + # This would store the communication in the database + # For now, just log it + logger.info( + f"Agent communication: {communication.from_agent_id} -> {communication.to_agent_id}: {communication.content}" + ) + + async def _handle_communication_timeout( + self, session: CoordinationSession, communication: AgentCommunication + ): + """Handle communication timeout""" + logger.warning( + f"Communication timeout in session {session.session_id}: {communication.communication_id}" + ) + + # Could implement retry logic or escalation here + + async def _synchronize_session(self, session: CoordinationSession): + """Synchronize coordination session state""" + # Check if any agents need help or coordination + # Update session status based on subtask progress + # Handle any conflicts or issues + pass + + async def _handle_coordination_failures(self, session: CoordinationSession): + """Handle failures in coordination""" + logger.warning( + f"Handling failures in coordination session {session.session_id}" + ) + + # Could implement retry logic, reassignment, or escalation + session.status = CoordinationStatus.FAILED + + async def _aggregate_coordination_results( + self, results: Dict[str, Dict[str, Any]] + ) -> Dict[str, Any]: + """Aggregate results from all coordinated subtasks""" + + successful_subtasks = [ + r for r in results.values() if r["status"] == "completed" + ] + + return { + "coordination_completed": True, + "total_subtasks": len(results), + "successful_subtasks": len(successful_subtasks), + "failed_subtasks": len(results) - len(successful_subtasks), + "results": results, + "completion_time": datetime.now().isoformat(), + } + + async def _update_agent_capabilities(self, agent_ids: List[str]): + """Update cached agent capabilities""" + for agent_id in agent_ids: + agent_data = await DatabaseManager.get_agent(agent_id) + if agent_data: + # Extract capabilities from agent configuration + tools = agent_data.get("config", {}).get("tools", []) + capabilities = [ + AgentCapability( + skill=tool, + proficiency=0.8, # Default proficiency + availability=agent_data.get("status") == "available", + current_load=0.5, # Default load + ) + for tool in tools + ] + self.agent_capabilities[agent_id] = capabilities + + async def _cleanup_session(self, session_id: str): + """Clean up completed coordination session""" + session = self.active_sessions.pop(session_id, None) + if session: + logger.info(f"Cleaned up coordination session {session_id}") + + +# Integration with existing TaskExecutionEngine +def integrate_multi_agent_coordination( + task_execution_engine: TaskExecutionEngine, +) -> MultiAgentCoordinator: + """Create and integrate multi-agent coordinator with task execution engine""" + coordinator = MultiAgentCoordinator(task_execution_engine) + + # Add coordination capabilities to task execution engine + task_execution_engine.multi_agent_coordinator = coordinator + + return coordinator diff --git a/services/orchestrator/task_execution_engine.py b/services/orchestrator/task_execution_engine.py index b11cd75..b2f3b3a 100644 --- a/services/orchestrator/task_execution_engine.py +++ b/services/orchestrator/task_execution_engine.py @@ -1,1306 +1,1306 @@ -""" -Task Execution Engine for FuzeAgent Autonomous Execution - -Orchestrates the autonomous execution of tasks by agents, managing: -- Task lifecycle and state transitions -- Sandbox creation and cleanup -- Git workflow automation -- Human-in-the-loop interactions -- Inter-agent communication -- Result aggregation - -This is the core component that ties together all autonomous execution components. -""" - -import asyncio -import json -import logging -import uuid -from dataclasses import dataclass -from datetime import datetime, timedelta -from enum import Enum -from typing import Any, Callable, Dict, List, Optional - -from .claude_code_wrapper import ClaudeCodeWrapper -from .claude_sdk_manager import ClaudeSDKManager, ClaudeSDKSession -from .context_enhancement_service import ContextEnhancementService -from .conversation_manager import ConversationManager, InteractionType -from .database import DatabaseManager, get_db_connection -from .file_operations_engine import FileOperationsEngine -from .git_workflow_manager import GitWorkflowManager -from .sandbox_manager import AgentSandboxManager, Sandbox -from .task_knowledge_extractor import TaskKnowledgeExtractor - -logger = logging.getLogger(__name__) - - -class TaskStatus(str, Enum): - PENDING = "pending" - ANALYZING = "analyzing" - SETTING_UP = "setting_up" - EXECUTING = "executing" - WAITING_FOR_HUMAN = "waiting_for_human" - REVIEWING = "reviewing" - COMMITTING = "committing" - COMPLETED = "completed" - FAILED = "failed" - CANCELLED = "cancelled" - - -class ExecutionStep(str, Enum): - ANALYZE_TASK = "analyze_task" - SETUP_SANDBOX = "setup_sandbox" - SETUP_GIT = "setup_git" - EXECUTE_ITERATION = "execute_iteration" - REVIEW_CHANGES = "review_changes" - COMMIT_CHANGES = "commit_changes" - HUMAN_INTERACTION = "human_interaction" - FINALIZE_TASK = "finalize_task" - CLEANUP = "cleanup" - - -@dataclass -class TaskIteration: - """Represents a single iteration of task execution""" - - iteration_number: int - step: ExecutionStep - started_at: datetime - completed_at: Optional[datetime] - input_data: Dict[str, Any] - output_data: Optional[Dict[str, Any]] - success: bool - error_message: Optional[str] - human_question: Optional[str] = None - human_response: Optional[str] = None - - -@dataclass -class ExecutionContext: - """Context for task execution""" - - task_id: str - agent_id: str - task_data: Dict[str, Any] - agent_data: Dict[str, Any] - sandbox: Optional[Sandbox] - git_manager: Optional[GitWorkflowManager] - claude_wrapper: Optional[ClaudeCodeWrapper] - current_iteration: int - iterations: List[TaskIteration] - status: TaskStatus - started_at: datetime - completed_at: Optional[datetime] - result: Optional[Dict[str, Any]] - error: Optional[str] - # New components for autonomous execution - file_operations_engine: Optional[FileOperationsEngine] = None - claude_sdk_manager: Optional[ClaudeSDKManager] = None - claude_session_id: Optional[str] = None - - -class TaskExecutionEngine: - """ - Orchestrates autonomous task execution by agents. - - Features: - - Task lifecycle management - - Sandbox and Git workflow integration - - Human-in-the-loop interactions - - Dependency handling - - Result aggregation - - Error recovery - """ - - def __init__( - self, - sandbox_manager: AgentSandboxManager, - knowledge_extractor: Optional[TaskKnowledgeExtractor] = None, - context_enhancer: Optional[ContextEnhancementService] = None, - ): - self.sandbox_manager = sandbox_manager - self.conversation_manager = ConversationManager() - self.active_executions: Dict[str, ExecutionContext] = {} - self.execution_callbacks: Dict[str, List[Callable]] = {} - self.running = False - self.worker_tasks: List[asyncio.Task] = [] - - # Knowledge management services - self.knowledge_extractor = knowledge_extractor - self.context_enhancer = context_enhancer - - # Initialize integrated components - self.file_operations_engines: Dict[str, FileOperationsEngine] = {} # Per task - self.claude_sdk_managers: Dict[str, ClaudeSDKManager] = {} # Per task - - # Configuration - self.max_iterations = 50 - self.iteration_timeout = 3600 # 1 hour per iteration - self.human_response_timeout = 86400 # 24 hours for human response - - async def start(self): - """Start the execution engine""" - logger.info("Starting TaskExecutionEngine") - self.running = True - - # Start worker tasks - self.worker_tasks = [ - asyncio.create_task(self._execution_worker()), - asyncio.create_task(self._monitoring_worker()), - asyncio.create_task(self._cleanup_worker()), - ] - - logger.info("TaskExecutionEngine started") - - async def stop(self): - """Stop the execution engine""" - logger.info("Stopping TaskExecutionEngine") - self.running = False - - # Cancel worker tasks - for task in self.worker_tasks: - task.cancel() - - try: - await asyncio.gather(*self.worker_tasks, return_exceptions=True) - except Exception as e: - logger.error(f"Error stopping worker tasks: {e}") - - # Clean up active executions - for execution_id in list(self.active_executions.keys()): - try: - await self._cleanup_execution(execution_id) - except Exception as e: - logger.error(f"Error cleaning up execution {execution_id}: {e}") - - logger.info("TaskExecutionEngine stopped") - - async def start_task_execution(self, task_id: str) -> Dict[str, Any]: - """ - Start autonomous execution of a task. - Returns execution status and context. - """ - logger.info(f"Starting task execution: {task_id}") - - try: - # Get task data - task_data = await self._get_task_data(task_id) - if not task_data: - raise ValueError(f"Task {task_id} not found") - - # Get agent data - agent_id = task_data.get("assigned_to") - if not agent_id: - raise ValueError(f"Task {task_id} has no assigned agent") - - agent_data = await self._get_agent_data(agent_id) - if not agent_data: - raise ValueError(f"Agent {agent_id} not found") - - # Create execution context - execution_context = ExecutionContext( - task_id=task_id, - agent_id=agent_id, - task_data=task_data, - agent_data=agent_data, - sandbox=None, - git_manager=None, - claude_wrapper=None, - current_iteration=0, - iterations=[], - status=TaskStatus.PENDING, - started_at=datetime.now(), - completed_at=None, - result=None, - error=None, - ) - - # Store execution context - self.active_executions[task_id] = execution_context - - # Update task status in database - await self._update_task_status(task_id, TaskStatus.PENDING) - - logger.info(f"✅ Task execution started: {task_id}") - return { - "task_id": task_id, - "status": TaskStatus.PENDING.value, - "execution_started": True, - "agent_id": agent_id, - } - - except Exception as e: - logger.error(f"❌ Failed to start task execution {task_id}: {e}") - await self._update_task_status(task_id, TaskStatus.FAILED, error=str(e)) - raise - - async def get_execution_status(self, task_id: str) -> Dict[str, Any]: - """Get detailed execution status for a task""" - - execution = self.active_executions.get(task_id) - if not execution: - # Check database for completed/failed tasks - task_data = await self._get_task_data(task_id) - if task_data: - return { - "task_id": task_id, - "status": task_data.get("status", "unknown"), - "result": task_data.get("result"), - "active_execution": False, - } - else: - return {"task_id": task_id, "status": "not_found"} - - return { - "task_id": task_id, - "status": execution.status.value, - "agent_id": execution.agent_id, - "current_iteration": execution.current_iteration, - "iterations_count": len(execution.iterations), - "started_at": execution.started_at.isoformat(), - "completed_at": ( - execution.completed_at.isoformat() if execution.completed_at else None - ), - "sandbox_id": execution.sandbox.sandbox_id if execution.sandbox else None, - "git_branch": ( - execution.git_manager.feature_branch if execution.git_manager else None - ), - "result": execution.result, - "error": execution.error, - "active_execution": True, - } - - async def get_task_iterations(self, task_id: str) -> List[Dict[str, Any]]: - """Get iteration history for a task""" - - execution = self.active_executions.get(task_id) - if execution: - iterations = execution.iterations - else: - # Get from database - iterations = await self._get_task_iterations_from_db(task_id) - - return [ - { - "iteration_number": it.iteration_number, - "step": it.step.value if hasattr(it.step, "value") else str(it.step), - "started_at": it.started_at.isoformat(), - "completed_at": ( - it.completed_at.isoformat() if it.completed_at else None - ), - "success": it.success, - "error_message": it.error_message, - "human_question": it.human_question, - "human_response": it.human_response, - "input_data": it.input_data, - "output_data": it.output_data, - } - for it in iterations - ] - - async def cancel_task_execution(self, task_id: str) -> bool: - """Cancel a running task execution""" - - execution = self.active_executions.get(task_id) - if not execution: - return False - - execution.status = TaskStatus.CANCELLED - execution.completed_at = datetime.now() - execution.error = "Task cancelled by user" - - # Update database - await self._update_task_status( - task_id, TaskStatus.CANCELLED, error="Task cancelled by user" - ) - - # Schedule cleanup - asyncio.create_task(self._cleanup_execution(task_id)) - - logger.info(f"Task execution cancelled: {task_id}") - return True - - # Private methods for execution workflow - - async def _execution_worker(self): - """Main execution worker that processes pending tasks""" - while self.running: - try: - # Find tasks ready for execution - pending_tasks = [ - task_id - for task_id, execution in self.active_executions.items() - if execution.status in [TaskStatus.PENDING, TaskStatus.EXECUTING] - ] - - # Process each pending task - for task_id in pending_tasks: - try: - await self._process_task_execution(task_id) - except Exception as e: - logger.error(f"Error processing task {task_id}: {e}") - await self._handle_execution_error(task_id, str(e)) - - # Sleep between iterations - await asyncio.sleep(5) - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Error in execution worker: {e}") - await asyncio.sleep(10) - - async def _process_task_execution(self, task_id: str): - """Process a single task execution step""" - execution = self.active_executions.get(task_id) - if not execution: - return - - # Skip if waiting for human or in terminal state - if execution.status in [ - TaskStatus.WAITING_FOR_HUMAN, - TaskStatus.COMPLETED, - TaskStatus.FAILED, - TaskStatus.CANCELLED, - ]: - return - - # Determine next step - next_step = self._determine_next_step(execution) - if not next_step: - return - - # Execute the step - try: - await self._execute_step(execution, next_step) - except Exception as e: - logger.error(f"Error executing step {next_step} for task {task_id}: {e}") - await self._handle_execution_error(task_id, str(e)) - - def _determine_next_step( - self, execution: ExecutionContext - ) -> Optional[ExecutionStep]: - """Determine the next execution step""" - - if execution.status == TaskStatus.PENDING: - return ExecutionStep.ANALYZE_TASK - - if not execution.iterations: - return ExecutionStep.ANALYZE_TASK - - last_iteration = execution.iterations[-1] - - # Continue based on last completed step - if last_iteration.step == ExecutionStep.ANALYZE_TASK and last_iteration.success: - return ExecutionStep.SETUP_SANDBOX - elif ( - last_iteration.step == ExecutionStep.SETUP_SANDBOX - and last_iteration.success - ): - return ExecutionStep.SETUP_GIT - elif last_iteration.step == ExecutionStep.SETUP_GIT and last_iteration.success: - return ExecutionStep.EXECUTE_ITERATION - elif ( - last_iteration.step == ExecutionStep.EXECUTE_ITERATION - and last_iteration.success - ): - # Check if we need human input - if last_iteration.human_question: - return ExecutionStep.HUMAN_INTERACTION - else: - return ExecutionStep.REVIEW_CHANGES - elif ( - last_iteration.step == ExecutionStep.HUMAN_INTERACTION - and last_iteration.human_response - ): - return ExecutionStep.EXECUTE_ITERATION - elif ( - last_iteration.step == ExecutionStep.REVIEW_CHANGES - and last_iteration.success - ): - return ExecutionStep.COMMIT_CHANGES - elif ( - last_iteration.step == ExecutionStep.COMMIT_CHANGES - and last_iteration.success - ): - # Check if task is complete - if self._is_task_complete(execution): - return ExecutionStep.FINALIZE_TASK - else: - return ExecutionStep.EXECUTE_ITERATION - - return None - - async def _execute_step(self, execution: ExecutionContext, step: ExecutionStep): - """Execute a specific step""" - - iteration = TaskIteration( - iteration_number=execution.current_iteration + 1, - step=step, - started_at=datetime.now(), - completed_at=None, - input_data={}, - output_data=None, - success=False, - error_message=None, - ) - - execution.iterations.append(iteration) - execution.current_iteration += 1 - - try: - if step == ExecutionStep.ANALYZE_TASK: - await self._step_analyze_task(execution, iteration) - elif step == ExecutionStep.SETUP_SANDBOX: - await self._step_setup_sandbox(execution, iteration) - elif step == ExecutionStep.SETUP_GIT: - await self._step_setup_git(execution, iteration) - elif step == ExecutionStep.EXECUTE_ITERATION: - await self._step_execute_iteration(execution, iteration) - elif step == ExecutionStep.REVIEW_CHANGES: - await self._step_review_changes(execution, iteration) - elif step == ExecutionStep.COMMIT_CHANGES: - await self._step_commit_changes(execution, iteration) - elif step == ExecutionStep.HUMAN_INTERACTION: - await self._step_human_interaction(execution, iteration) - elif step == ExecutionStep.FINALIZE_TASK: - await self._step_finalize_task(execution, iteration) - - iteration.completed_at = datetime.now() - iteration.success = True - - except Exception as e: - iteration.completed_at = datetime.now() - iteration.success = False - iteration.error_message = str(e) - raise - - finally: - # Store iteration in database - await self._store_task_iteration(execution.task_id, iteration) - - async def _step_analyze_task( - self, execution: ExecutionContext, iteration: TaskIteration - ): - """Analyze the task and prepare execution plan""" - execution.status = TaskStatus.ANALYZING - await self._update_task_status(execution.task_id, TaskStatus.ANALYZING) - - # Analyze task requirements - task_description = execution.task_data.get("description", "") - task_title = execution.task_data.get("title", "") - - iteration.input_data = { - "task_title": task_title, - "task_description": task_description, - "agent_type": execution.agent_data.get("type"), - "agent_role": execution.agent_data.get("role"), - } - - # Enhance context with organizational knowledge - enhanced_context = None - if self.context_enhancer: - try: - enhanced_context = await self.context_enhancer.enhance_agent_context( - agent_id=execution.agent_id, - task_data=execution.task_data, - base_context=iteration.input_data, - ) - logger.info( - f"Enhanced context for task {execution.task_id}: " - f"{len(enhanced_context.organizational_knowledge)} org + " - f"{len(enhanced_context.team_knowledge)} team + " - f"{len(enhanced_context.similar_task_insights)} similar task insights" - ) - except Exception as e: - logger.error( - f"Error enhancing context for task {execution.task_id}: {e}" - ) - - # Simple analysis for now - in a full implementation this would use AI - iteration.output_data = { - "analysis_complete": True, - "requires_sandbox": execution.agent_data.get("type") == "developer", - "requires_git": bool( - execution.agent_data.get("repository_settings", {}).get( - "repository_url" - ) - ), - "enhanced_context": enhanced_context, - "estimated_complexity": "medium", - "estimated_iterations": 5, - } - - logger.info(f"Task analysis complete for {execution.task_id}") - - async def _step_setup_sandbox( - self, execution: ExecutionContext, iteration: TaskIteration - ): - """Set up sandbox environment for the agent""" - execution.status = TaskStatus.SETTING_UP - await self._update_task_status(execution.task_id, TaskStatus.SETTING_UP) - - agent_template = execution.agent_data.get("template_id", "python_developer") - repository_settings = execution.agent_data.get("repository_settings", {}) - sandbox_settings = execution.agent_data.get("sandbox_settings", {}) - - # Create sandbox - sandbox = await self.sandbox_manager.create_sandbox( - agent_id=execution.agent_id, - task_id=execution.task_id, - agent_template=agent_template, - repository_settings=repository_settings, - custom_settings=sandbox_settings, - ) - - execution.sandbox = sandbox - - iteration.input_data = { - "agent_template": agent_template, - "repository_settings": repository_settings, - "sandbox_settings": sandbox_settings, - } - - iteration.output_data = { - "sandbox_id": sandbox.sandbox_id, - "workspace_path": sandbox.workspace_path, - "container_id": sandbox.container_id, - } - - logger.info( - f"Sandbox setup complete for {execution.task_id}: {sandbox.sandbox_id}" - ) - - async def _step_setup_git( - self, execution: ExecutionContext, iteration: TaskIteration - ): - """Set up Git workflow for the task""" - repository_settings = execution.agent_data.get("repository_settings", {}) - - if not repository_settings.get("repository_url"): - # Skip Git setup if no repository - iteration.output_data = { - "git_setup": "skipped", - "reason": "no_repository_configured", - } - return - - # Create Git workflow manager - git_manager = GitWorkflowManager( - agent_id=execution.agent_id, - task_id=execution.task_id, - repo_settings=repository_settings, - ) - - # Setup workspace - feature_branch = await git_manager.setup_workspace() - - execution.git_manager = git_manager - - # Create enhanced Claude wrapper with Git context and conversation tracking - execution.claude_wrapper = ClaudeCodeWrapper( - workspace_path=git_manager.workspace_path, - git_manager=git_manager, - agent_id=execution.agent_id, - task_id=execution.task_id, - conversation_manager=self.conversation_manager, - ) - - # Initialize File Operations Engine - file_ops_engine = FileOperationsEngine(git_manager.workspace_path) - execution.file_operations_engine = file_ops_engine - self.file_operations_engines[execution.task_id] = file_ops_engine - - # Initialize Claude SDK Manager - claude_sdk_manager = ClaudeSDKManager( - file_operations_engine=file_ops_engine, - conversation_manager=self.conversation_manager, - ) - execution.claude_sdk_manager = claude_sdk_manager - self.claude_sdk_managers[execution.task_id] = claude_sdk_manager - - # Start conversation session - await execution.claude_wrapper.start_conversation_session( - execution.sandbox.sandbox_id - ) - - iteration.input_data = { - "repository_url": repository_settings.get("repository_url"), - "default_branch": repository_settings.get("default_branch", "main"), - } - - iteration.output_data = { - "git_setup": "complete", - "feature_branch": feature_branch, - "workspace_path": git_manager.workspace_path, - } - - logger.info(f"Git setup complete for {execution.task_id}: {feature_branch}") - - async def _step_execute_iteration( - self, execution: ExecutionContext, iteration: TaskIteration - ): - """Execute a development iteration using Claude SDK Manager""" - execution.status = TaskStatus.EXECUTING - await self._update_task_status(execution.task_id, TaskStatus.EXECUTING) - - task_description = execution.task_data.get("description", "") - task_title = execution.task_data.get("title", "") - - iteration.input_data = { - "task_description": task_description, - "task_title": task_title, - "iteration_number": iteration.iteration_number, - "workspace_path": ( - execution.sandbox.workspace_path if execution.sandbox else None - ), - } - - try: - # Start Claude SDK session if not already running - if not execution.claude_session_id and execution.claude_sdk_manager: - # Build comprehensive task context - context_info = "" - if execution.git_manager: - context_info += f"\nRepository: {execution.git_manager.repo_url}" - context_info += f"\nBranch: {execution.git_manager.feature_branch}" - if iteration.iteration_number > 1: - context_info += f"\nIteration: {iteration.iteration_number} of ongoing development" - - # Construct task prompt - task_prompt = f""" -Task: {task_title} - -Description: {task_description} - -Context: {context_info} - -Please analyze the codebase, understand the requirements, and implement the necessary changes. -Work incrementally and ask for clarification if needed. -""" - - # Start Claude SDK session - execution.claude_session_id = ( - await execution.claude_sdk_manager.start_session( - task_id=execution.task_id, - agent_id=execution.agent_id, - workspace_path=( - execution.git_manager.workspace_path - if execution.git_manager - else execution.sandbox.workspace_path - ), - task_description=task_prompt, - additional_context=context_info, - ) - ) - - logger.info( - f"Started Claude SDK session: {execution.claude_session_id}" - ) - - # Register interaction callback for human-in-the-loop - if execution.claude_sdk_manager and execution.claude_session_id: - execution.claude_sdk_manager.register_interaction_callback( - execution.claude_session_id, - lambda session, interaction: self._handle_claude_interaction( - execution, session, interaction - ), - ) - - # Monitor session status - session_status = ( - await execution.claude_sdk_manager.get_session_status( - execution.claude_session_id - ) - if execution.claude_session_id - else None - ) - - if session_status: - current_interaction = session_status.get("current_interaction") - if current_interaction: - # Claude is waiting for human input - interaction_type = current_interaction.get("type") - if interaction_type in ["user_input", "confirmation"]: - iteration.human_question = current_interaction.get("prompt") - execution.status = TaskStatus.WAITING_FOR_HUMAN - await self._update_task_status( - execution.task_id, TaskStatus.WAITING_FOR_HUMAN - ) - - iteration.output_data = { - "claude_session_status": session_status.get("state"), - "human_interaction_required": True, - "interaction_type": interaction_type, - "human_question_asked": True, - } - - logger.info( - f"Claude SDK requesting human input for task {execution.task_id}" - ) - return - - elif interaction_type == "file_approval": - # File operations pending approval - batch_id = current_interaction.get("metadata", {}).get( - "batch_id" - ) - if batch_id and execution.file_operations_engine: - # Get diff preview for human review - diffs = await execution.file_operations_engine.get_file_diff_preview( - batch_id - ) - - iteration.human_question = f"""Claude wants to make the following file changes: - -{current_interaction.get('prompt')} - -File changes preview: -{self._format_diffs_for_human(diffs)} - -Approve these changes? (yes/no)""" - - execution.status = TaskStatus.WAITING_FOR_HUMAN - await self._update_task_status( - execution.task_id, TaskStatus.WAITING_FOR_HUMAN - ) - - iteration.output_data = { - "claude_session_status": session_status.get("state"), - "file_approval_required": True, - "batch_id": batch_id, - "file_changes_preview": diffs, - "human_question_asked": True, - } - - logger.info( - f"Claude SDK requesting file approval for task {execution.task_id}" - ) - return - - # No interaction needed - continue execution - iteration.output_data = { - "claude_session_status": session_status.get("state"), - "session_active": True, - "iteration_completed": True, - "workspace_path": session_status.get("workspace_path"), - } - - # Check if session completed - if session_status.get("state") in ["completed", "terminated"]: - iteration.output_data["development_complete"] = True - - else: - # No session - this shouldn't happen but handle gracefully - iteration.output_data = { - "error": "No Claude SDK session available", - "development_complete": False, - } - - except Exception as e: - logger.error(f"Error in Claude SDK iteration: {e}") - iteration.output_data = {"error": str(e), "development_complete": False} - raise - - logger.info( - f"Development iteration {iteration.iteration_number} processed for {execution.task_id}" - ) - - async def _step_review_changes( - self, execution: ExecutionContext, iteration: TaskIteration - ): - """Review the changes made in the iteration""" - execution.status = TaskStatus.REVIEWING - await self._update_task_status(execution.task_id, TaskStatus.REVIEWING) - - # Review changes - in full implementation this would: - # 1. Run linting and type checking - # 2. Run tests - # 3. Check code quality - # 4. Validate against requirements - - iteration.output_data = { - "review_passed": True, - "issues_found": [], - "tests_passed": True, - "code_quality_score": 85, - } - - logger.info(f"Code review complete for {execution.task_id}") - - async def _step_commit_changes( - self, execution: ExecutionContext, iteration: TaskIteration - ): - """Commit changes to Git""" - execution.status = TaskStatus.COMMITTING - await self._update_task_status(execution.task_id, TaskStatus.COMMITTING) - - if not execution.git_manager: - iteration.output_data = {"commit": "skipped", "reason": "no_git_manager"} - return - - # Commit changes - commit_message = f"Iteration {iteration.iteration_number}: {execution.task_data.get('title', 'Task update')}" - commit_hash = await execution.git_manager.commit_changes( - message=commit_message, iteration_number=iteration.iteration_number - ) - - iteration.output_data = { - "commit_hash": commit_hash, - "commit_message": commit_message, - "branch": execution.git_manager.feature_branch, - } - - logger.info(f"Changes committed for {execution.task_id}: {commit_hash}") - - async def _step_human_interaction( - self, execution: ExecutionContext, iteration: TaskIteration - ): - """Handle human interaction step""" - # This step waits for human response - the actual waiting is handled - # by the status being WAITING_FOR_HUMAN - iteration.output_data = { - "human_interaction": "waiting_for_response", - "question": iteration.human_question, - } - - async def _step_finalize_task( - self, execution: ExecutionContext, iteration: TaskIteration - ): - """Finalize the task execution""" - execution.status = TaskStatus.COMPLETED - execution.completed_at = datetime.now() - - # Create pull request if Git is configured - pr_url = None - if execution.git_manager: - try: - pr = await execution.git_manager.create_pull_request( - title=f"🤖 {execution.task_data.get('title', 'Task completion')}", - description=f"Autonomous completion of task: {execution.task_data.get('description', '')}", - ) - pr_url = pr.url if pr else None - except Exception as e: - logger.error(f"Failed to create PR for {execution.task_id}: {e}") - - # Prepare final result - execution.result = { - "status": "completed", - "iterations": len(execution.iterations), - "started_at": execution.started_at.isoformat(), - "completed_at": execution.completed_at.isoformat(), - "pull_request_url": pr_url, - "sandbox_id": execution.sandbox.sandbox_id if execution.sandbox else None, - "git_branch": ( - execution.git_manager.feature_branch if execution.git_manager else None - ), - } - - # Update database - await self._update_task_status( - execution.task_id, TaskStatus.COMPLETED, result=execution.result - ) - - iteration.output_data = execution.result - - # End conversation session - if execution.claude_wrapper: - try: - await execution.claude_wrapper.end_conversation_session() - except Exception as e: - logger.error(f"Error ending conversation session: {e}") - - # Store final performance metrics - await self._store_completion_metrics(execution) - - # Extract knowledge from completed task - if self.knowledge_extractor: - try: - extracted_knowledge_ids = ( - await self.knowledge_extractor.extract_knowledge_from_task( - task_id=execution.task_id, - agent_id=execution.agent_id, - execution_result=execution.result, - ) - ) - if extracted_knowledge_ids: - logger.info( - f"Extracted {len(extracted_knowledge_ids)} knowledge items from task {execution.task_id}" - ) - except Exception as e: - logger.error( - f"Error extracting knowledge from task {execution.task_id}: {e}" - ) - - logger.info(f"✅ Task execution completed: {execution.task_id}") - - # End Claude SDK session - if execution.claude_sdk_manager and execution.claude_session_id: - try: - await execution.claude_sdk_manager.terminate_session( - execution.claude_session_id - ) - except Exception as e: - logger.error(f"Error terminating Claude SDK session: {e}") - - # Schedule cleanup - asyncio.create_task(self._cleanup_execution(execution.task_id)) - - def _is_task_complete(self, execution: ExecutionContext) -> bool: - """Check if the task is complete""" - # Simple completion check - in full implementation this would be more sophisticated - return execution.current_iteration >= 3 - - async def _handle_execution_error(self, task_id: str, error_message: str): - """Handle execution error""" - execution = self.active_executions.get(task_id) - if execution: - execution.status = TaskStatus.FAILED - execution.completed_at = datetime.now() - execution.error = error_message - - await self._update_task_status(task_id, TaskStatus.FAILED, error=error_message) - - # Schedule cleanup - asyncio.create_task(self._cleanup_execution(task_id)) - - logger.error(f"Task execution failed: {task_id} - {error_message}") - - async def _cleanup_execution(self, task_id: str): - """Clean up execution resources""" - execution = self.active_executions.pop(task_id, None) - if not execution: - return - - # Cleanup Claude SDK session - if execution.claude_sdk_manager and execution.claude_session_id: - try: - await execution.claude_sdk_manager.terminate_session( - execution.claude_session_id - ) - except Exception as e: - logger.error(f"Error terminating Claude SDK session for {task_id}: {e}") - - # Remove engines from tracking - self.file_operations_engines.pop(task_id, None) - self.claude_sdk_managers.pop(task_id, None) - - # Cleanup sandbox - if execution.sandbox: - try: - await self.sandbox_manager.destroy_sandbox(execution.sandbox.sandbox_id) - except Exception as e: - logger.error(f"Error destroying sandbox for {task_id}: {e}") - - # Cleanup Git workspace (optional - might want to keep for review) - if execution.git_manager and execution.status != TaskStatus.COMPLETED: - try: - await execution.git_manager.cleanup_workspace() - except Exception as e: - logger.error(f"Error cleaning up Git workspace for {task_id}: {e}") - - logger.info(f"Execution cleanup complete: {task_id}") - - async def _monitoring_worker(self): - """Monitor execution health and timeouts""" - while self.running: - try: - current_time = datetime.now() - - for task_id, execution in list(self.active_executions.items()): - # Check for timeouts - if execution.status == TaskStatus.WAITING_FOR_HUMAN: - # Check human response timeout - last_iteration = ( - execution.iterations[-1] if execution.iterations else None - ) - if last_iteration and last_iteration.started_at: - time_waiting = current_time - last_iteration.started_at - if time_waiting > timedelta( - seconds=self.human_response_timeout - ): - await self._handle_execution_error( - task_id, - "Human response timeout - no response received within 24 hours", - ) - else: - # Check general execution timeout - execution_time = current_time - execution.started_at - if execution_time > timedelta( - seconds=self.iteration_timeout * self.max_iterations - ): - await self._handle_execution_error( - task_id, - f"Execution timeout - exceeded maximum time limit", - ) - - # Check iteration limits - if execution.current_iteration > self.max_iterations: - await self._handle_execution_error( - task_id, - f"Maximum iterations exceeded ({self.max_iterations})", - ) - - await asyncio.sleep(60) # Check every minute - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Error in monitoring worker: {e}") - await asyncio.sleep(60) - - async def _cleanup_worker(self): - """Clean up completed executions periodically""" - while self.running: - try: - current_time = datetime.now() - cutoff_time = current_time - timedelta( - hours=1 - ) # Keep completed tasks for 1 hour - - completed_tasks = [ - task_id - for task_id, execution in self.active_executions.items() - if execution.status - in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED] - and execution.completed_at - and execution.completed_at < cutoff_time - ] - - for task_id in completed_tasks: - await self._cleanup_execution(task_id) - - await asyncio.sleep(3600) # Run every hour - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Error in cleanup worker: {e}") - await asyncio.sleep(3600) - - # Database operations - - async def _get_task_data(self, task_id: str) -> Optional[Dict[str, Any]]: - """Get task data from database""" - async with get_db_connection() as conn: - row = await conn.fetchrow("SELECT * FROM tasks WHERE id = $1", task_id) - return dict(row) if row else None - - async def _get_agent_data(self, agent_id: str) -> Optional[Dict[str, Any]]: - """Get agent data from database""" - return await DatabaseManager.get_agent(agent_id) - - async def _update_task_status( - self, - task_id: str, - status: TaskStatus, - result: Optional[Dict[str, Any]] = None, - error: Optional[str] = None, - ): - """Update task status in database""" - await DatabaseManager.update_task_status( - task_id=task_id, status=status.value, result=result - ) - - async def _store_task_iteration(self, task_id: str, iteration: TaskIteration): - """Store task iteration in database""" - async with get_db_connection() as conn: - await conn.execute( - """ - INSERT INTO task_iterations ( - id, task_id, iteration_number, step, started_at, completed_at, - input_data, output_data, success, error_message, - human_question, human_response - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) - """, - str(uuid.uuid4()), - task_id, - iteration.iteration_number, - iteration.step.value, - iteration.started_at, - iteration.completed_at, - json.dumps(iteration.input_data), - json.dumps(iteration.output_data) if iteration.output_data else None, - iteration.success, - iteration.error_message, - iteration.human_question, - iteration.human_response, - ) - - async def _get_task_iterations_from_db(self, task_id: str) -> List[TaskIteration]: - """Get task iterations from database""" - async with get_db_connection() as conn: - rows = await conn.fetch( - """ - SELECT * FROM task_iterations - WHERE task_id = $1 - ORDER BY iteration_number - """, - task_id, - ) - - iterations = [] - for row in rows: - iterations.append( - TaskIteration( - iteration_number=row["iteration_number"], - step=ExecutionStep(row["step"]), - started_at=row["started_at"], - completed_at=row["completed_at"], - input_data=( - json.loads(row["input_data"]) if row["input_data"] else {} - ), - output_data=( - json.loads(row["output_data"]) - if row["output_data"] - else None - ), - success=row["success"], - error_message=row["error_message"], - human_question=row["human_question"], - human_response=row["human_response"], - ) - ) - - return iterations - - async def _store_completion_metrics(self, execution: ExecutionContext): - """Store performance metrics for completed task""" - try: - if not execution.completed_at or not execution.started_at: - return - - # Calculate execution time - execution_time = ( - execution.completed_at - execution.started_at - ).total_seconds() / 60 # minutes - - # Store metrics - await self.conversation_manager.store_performance_metric( - agent_id=execution.agent_id, - task_id=execution.task_id, - metric_type="execution_time_minutes", - metric_value=execution_time, - metric_unit="minutes", - ) - - await self.conversation_manager.store_performance_metric( - agent_id=execution.agent_id, - task_id=execution.task_id, - metric_type="iterations_to_completion", - metric_value=float(execution.current_iteration), - metric_unit="iterations", - ) - - # Store success/failure metric - success_value = 1.0 if execution.status == TaskStatus.COMPLETED else 0.0 - await self.conversation_manager.store_performance_metric( - agent_id=execution.agent_id, - task_id=execution.task_id, - metric_type="task_success_rate", - metric_value=success_value, - metric_unit="boolean", - ) - - except Exception as e: - logger.error(f"Error storing completion metrics: {e}") - - async def _handle_claude_interaction( - self, execution: ExecutionContext, session: ClaudeSDKSession, interaction - ): - """Handle interaction from Claude SDK""" - logger.info( - f"Claude interaction for task {execution.task_id}: {interaction.interaction_type}" - ) - - # This will be processed in the next iteration of _step_execute_iteration - # The interaction handling is done there to maintain the execution flow - - def _format_diffs_for_human(self, diffs: Dict[str, str]) -> str: - """Format file diffs for human review""" - if not diffs: - return "No file changes detected." - - formatted = [] - for file_path, diff in diffs.items(): - formatted.append(f"\n--- {file_path} ---") - formatted.append(diff[:1000] + "..." if len(diff) > 1000 else diff) - - return "\n".join(formatted) - - async def handle_human_response(self, task_id: str, response: str) -> bool: - """Enhanced human response handler that integrates with Claude SDK""" - - execution = self.active_executions.get(task_id) - if not execution or execution.status != TaskStatus.WAITING_FOR_HUMAN: - return False - - # Find the current iteration waiting for human response - current_iteration = None - for iteration in reversed(execution.iterations): - if iteration.human_question and not iteration.human_response: - current_iteration = iteration - break - - if current_iteration: - current_iteration.human_response = response - execution.status = TaskStatus.EXECUTING - - # Handle different types of responses - output_data = current_iteration.output_data or {} - - if output_data.get("file_approval_required"): - # Handle file approval - batch_id = output_data.get("batch_id") - approved = response.lower().strip() in [ - "yes", - "y", - "approve", - "approved", - "true", - ] - - if ( - batch_id - and execution.claude_sdk_manager - and execution.claude_session_id - ): - success = ( - await execution.claude_sdk_manager.approve_file_operations( - execution.claude_session_id, batch_id, approved - ) - ) - - if success: - logger.info( - f"File operations {'approved' if approved else 'rejected'} for task {task_id}" - ) - else: - logger.error( - f"Failed to process file approval for task {task_id}" - ) - - else: - # Handle general user input - if execution.claude_sdk_manager and execution.claude_session_id: - success = await execution.claude_sdk_manager.send_input( - execution.claude_session_id, response - ) - - if success: - logger.info( - f"Human response sent to Claude SDK for task {task_id}" - ) - else: - logger.error( - f"Failed to send human response to Claude SDK for task {task_id}" - ) - - # Update database - await self._store_task_iteration(task_id, current_iteration) - await self._update_task_status(task_id, TaskStatus.EXECUTING) - - logger.info(f"Human response processed for task {task_id}") - return True - - return False +""" +Task Execution Engine for FuzeAgent Autonomous Execution + +Orchestrates the autonomous execution of tasks by agents, managing: +- Task lifecycle and state transitions +- Sandbox creation and cleanup +- Git workflow automation +- Human-in-the-loop interactions +- Inter-agent communication +- Result aggregation + +This is the core component that ties together all autonomous execution components. +""" + +import asyncio +import json +import logging +import uuid +from dataclasses import dataclass +from datetime import datetime, timedelta +from enum import Enum +from typing import Any, Callable, Dict, List, Optional + +from .claude_code_wrapper import ClaudeCodeWrapper +from .claude_sdk_manager import ClaudeSDKManager, ClaudeSDKSession +from .context_enhancement_service import ContextEnhancementService +from .conversation_manager import ConversationManager, InteractionType +from .database import DatabaseManager, get_db_connection +from .file_operations_engine import FileOperationsEngine +from .git_workflow_manager import GitWorkflowManager +from .sandbox_manager import AgentSandboxManager, Sandbox +from .task_knowledge_extractor import TaskKnowledgeExtractor + +logger = logging.getLogger(__name__) + + +class TaskStatus(str, Enum): + PENDING = "pending" + ANALYZING = "analyzing" + SETTING_UP = "setting_up" + EXECUTING = "executing" + WAITING_FOR_HUMAN = "waiting_for_human" + REVIEWING = "reviewing" + COMMITTING = "committing" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +class ExecutionStep(str, Enum): + ANALYZE_TASK = "analyze_task" + SETUP_SANDBOX = "setup_sandbox" + SETUP_GIT = "setup_git" + EXECUTE_ITERATION = "execute_iteration" + REVIEW_CHANGES = "review_changes" + COMMIT_CHANGES = "commit_changes" + HUMAN_INTERACTION = "human_interaction" + FINALIZE_TASK = "finalize_task" + CLEANUP = "cleanup" + + +@dataclass +class TaskIteration: + """Represents a single iteration of task execution""" + + iteration_number: int + step: ExecutionStep + started_at: datetime + completed_at: Optional[datetime] + input_data: Dict[str, Any] + output_data: Optional[Dict[str, Any]] + success: bool + error_message: Optional[str] + human_question: Optional[str] = None + human_response: Optional[str] = None + + +@dataclass +class ExecutionContext: + """Context for task execution""" + + task_id: str + agent_id: str + task_data: Dict[str, Any] + agent_data: Dict[str, Any] + sandbox: Optional[Sandbox] + git_manager: Optional[GitWorkflowManager] + claude_wrapper: Optional[ClaudeCodeWrapper] + current_iteration: int + iterations: List[TaskIteration] + status: TaskStatus + started_at: datetime + completed_at: Optional[datetime] + result: Optional[Dict[str, Any]] + error: Optional[str] + # New components for autonomous execution + file_operations_engine: Optional[FileOperationsEngine] = None + claude_sdk_manager: Optional[ClaudeSDKManager] = None + claude_session_id: Optional[str] = None + + +class TaskExecutionEngine: + """ + Orchestrates autonomous task execution by agents. + + Features: + - Task lifecycle management + - Sandbox and Git workflow integration + - Human-in-the-loop interactions + - Dependency handling + - Result aggregation + - Error recovery + """ + + def __init__( + self, + sandbox_manager: AgentSandboxManager, + knowledge_extractor: Optional[TaskKnowledgeExtractor] = None, + context_enhancer: Optional[ContextEnhancementService] = None, + ): + self.sandbox_manager = sandbox_manager + self.conversation_manager = ConversationManager() + self.active_executions: Dict[str, ExecutionContext] = {} + self.execution_callbacks: Dict[str, List[Callable]] = {} + self.running = False + self.worker_tasks: List[asyncio.Task] = [] + + # Knowledge management services + self.knowledge_extractor = knowledge_extractor + self.context_enhancer = context_enhancer + + # Initialize integrated components + self.file_operations_engines: Dict[str, FileOperationsEngine] = {} # Per task + self.claude_sdk_managers: Dict[str, ClaudeSDKManager] = {} # Per task + + # Configuration + self.max_iterations = 50 + self.iteration_timeout = 3600 # 1 hour per iteration + self.human_response_timeout = 86400 # 24 hours for human response + + async def start(self): + """Start the execution engine""" + logger.info("Starting TaskExecutionEngine") + self.running = True + + # Start worker tasks + self.worker_tasks = [ + asyncio.create_task(self._execution_worker()), + asyncio.create_task(self._monitoring_worker()), + asyncio.create_task(self._cleanup_worker()), + ] + + logger.info("TaskExecutionEngine started") + + async def stop(self): + """Stop the execution engine""" + logger.info("Stopping TaskExecutionEngine") + self.running = False + + # Cancel worker tasks + for task in self.worker_tasks: + task.cancel() + + try: + await asyncio.gather(*self.worker_tasks, return_exceptions=True) + except Exception as e: + logger.error(f"Error stopping worker tasks: {e}") + + # Clean up active executions + for execution_id in list(self.active_executions.keys()): + try: + await self._cleanup_execution(execution_id) + except Exception as e: + logger.error(f"Error cleaning up execution {execution_id}: {e}") + + logger.info("TaskExecutionEngine stopped") + + async def start_task_execution(self, task_id: str) -> Dict[str, Any]: + """ + Start autonomous execution of a task. + Returns execution status and context. + """ + logger.info(f"Starting task execution: {task_id}") + + try: + # Get task data + task_data = await self._get_task_data(task_id) + if not task_data: + raise ValueError(f"Task {task_id} not found") + + # Get agent data + agent_id = task_data.get("assigned_to") + if not agent_id: + raise ValueError(f"Task {task_id} has no assigned agent") + + agent_data = await self._get_agent_data(agent_id) + if not agent_data: + raise ValueError(f"Agent {agent_id} not found") + + # Create execution context + execution_context = ExecutionContext( + task_id=task_id, + agent_id=agent_id, + task_data=task_data, + agent_data=agent_data, + sandbox=None, + git_manager=None, + claude_wrapper=None, + current_iteration=0, + iterations=[], + status=TaskStatus.PENDING, + started_at=datetime.now(), + completed_at=None, + result=None, + error=None, + ) + + # Store execution context + self.active_executions[task_id] = execution_context + + # Update task status in database + await self._update_task_status(task_id, TaskStatus.PENDING) + + logger.info(f"✅ Task execution started: {task_id}") + return { + "task_id": task_id, + "status": TaskStatus.PENDING.value, + "execution_started": True, + "agent_id": agent_id, + } + + except Exception as e: + logger.error(f"❌ Failed to start task execution {task_id}: {e}") + await self._update_task_status(task_id, TaskStatus.FAILED, error=str(e)) + raise + + async def get_execution_status(self, task_id: str) -> Dict[str, Any]: + """Get detailed execution status for a task""" + + execution = self.active_executions.get(task_id) + if not execution: + # Check database for completed/failed tasks + task_data = await self._get_task_data(task_id) + if task_data: + return { + "task_id": task_id, + "status": task_data.get("status", "unknown"), + "result": task_data.get("result"), + "active_execution": False, + } + else: + return {"task_id": task_id, "status": "not_found"} + + return { + "task_id": task_id, + "status": execution.status.value, + "agent_id": execution.agent_id, + "current_iteration": execution.current_iteration, + "iterations_count": len(execution.iterations), + "started_at": execution.started_at.isoformat(), + "completed_at": ( + execution.completed_at.isoformat() if execution.completed_at else None + ), + "sandbox_id": execution.sandbox.sandbox_id if execution.sandbox else None, + "git_branch": ( + execution.git_manager.feature_branch if execution.git_manager else None + ), + "result": execution.result, + "error": execution.error, + "active_execution": True, + } + + async def get_task_iterations(self, task_id: str) -> List[Dict[str, Any]]: + """Get iteration history for a task""" + + execution = self.active_executions.get(task_id) + if execution: + iterations = execution.iterations + else: + # Get from database + iterations = await self._get_task_iterations_from_db(task_id) + + return [ + { + "iteration_number": it.iteration_number, + "step": it.step.value if hasattr(it.step, "value") else str(it.step), + "started_at": it.started_at.isoformat(), + "completed_at": ( + it.completed_at.isoformat() if it.completed_at else None + ), + "success": it.success, + "error_message": it.error_message, + "human_question": it.human_question, + "human_response": it.human_response, + "input_data": it.input_data, + "output_data": it.output_data, + } + for it in iterations + ] + + async def cancel_task_execution(self, task_id: str) -> bool: + """Cancel a running task execution""" + + execution = self.active_executions.get(task_id) + if not execution: + return False + + execution.status = TaskStatus.CANCELLED + execution.completed_at = datetime.now() + execution.error = "Task cancelled by user" + + # Update database + await self._update_task_status( + task_id, TaskStatus.CANCELLED, error="Task cancelled by user" + ) + + # Schedule cleanup + asyncio.create_task(self._cleanup_execution(task_id)) + + logger.info(f"Task execution cancelled: {task_id}") + return True + + # Private methods for execution workflow + + async def _execution_worker(self): + """Main execution worker that processes pending tasks""" + while self.running: + try: + # Find tasks ready for execution + pending_tasks = [ + task_id + for task_id, execution in self.active_executions.items() + if execution.status in [TaskStatus.PENDING, TaskStatus.EXECUTING] + ] + + # Process each pending task + for task_id in pending_tasks: + try: + await self._process_task_execution(task_id) + except Exception as e: + logger.error(f"Error processing task {task_id}: {e}") + await self._handle_execution_error(task_id, str(e)) + + # Sleep between iterations + await asyncio.sleep(5) + + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"Error in execution worker: {e}") + await asyncio.sleep(10) + + async def _process_task_execution(self, task_id: str): + """Process a single task execution step""" + execution = self.active_executions.get(task_id) + if not execution: + return + + # Skip if waiting for human or in terminal state + if execution.status in [ + TaskStatus.WAITING_FOR_HUMAN, + TaskStatus.COMPLETED, + TaskStatus.FAILED, + TaskStatus.CANCELLED, + ]: + return + + # Determine next step + next_step = self._determine_next_step(execution) + if not next_step: + return + + # Execute the step + try: + await self._execute_step(execution, next_step) + except Exception as e: + logger.error(f"Error executing step {next_step} for task {task_id}: {e}") + await self._handle_execution_error(task_id, str(e)) + + def _determine_next_step( + self, execution: ExecutionContext + ) -> Optional[ExecutionStep]: + """Determine the next execution step""" + + if execution.status == TaskStatus.PENDING: + return ExecutionStep.ANALYZE_TASK + + if not execution.iterations: + return ExecutionStep.ANALYZE_TASK + + last_iteration = execution.iterations[-1] + + # Continue based on last completed step + if last_iteration.step == ExecutionStep.ANALYZE_TASK and last_iteration.success: + return ExecutionStep.SETUP_SANDBOX + elif ( + last_iteration.step == ExecutionStep.SETUP_SANDBOX + and last_iteration.success + ): + return ExecutionStep.SETUP_GIT + elif last_iteration.step == ExecutionStep.SETUP_GIT and last_iteration.success: + return ExecutionStep.EXECUTE_ITERATION + elif ( + last_iteration.step == ExecutionStep.EXECUTE_ITERATION + and last_iteration.success + ): + # Check if we need human input + if last_iteration.human_question: + return ExecutionStep.HUMAN_INTERACTION + else: + return ExecutionStep.REVIEW_CHANGES + elif ( + last_iteration.step == ExecutionStep.HUMAN_INTERACTION + and last_iteration.human_response + ): + return ExecutionStep.EXECUTE_ITERATION + elif ( + last_iteration.step == ExecutionStep.REVIEW_CHANGES + and last_iteration.success + ): + return ExecutionStep.COMMIT_CHANGES + elif ( + last_iteration.step == ExecutionStep.COMMIT_CHANGES + and last_iteration.success + ): + # Check if task is complete + if self._is_task_complete(execution): + return ExecutionStep.FINALIZE_TASK + else: + return ExecutionStep.EXECUTE_ITERATION + + return None + + async def _execute_step(self, execution: ExecutionContext, step: ExecutionStep): + """Execute a specific step""" + + iteration = TaskIteration( + iteration_number=execution.current_iteration + 1, + step=step, + started_at=datetime.now(), + completed_at=None, + input_data={}, + output_data=None, + success=False, + error_message=None, + ) + + execution.iterations.append(iteration) + execution.current_iteration += 1 + + try: + if step == ExecutionStep.ANALYZE_TASK: + await self._step_analyze_task(execution, iteration) + elif step == ExecutionStep.SETUP_SANDBOX: + await self._step_setup_sandbox(execution, iteration) + elif step == ExecutionStep.SETUP_GIT: + await self._step_setup_git(execution, iteration) + elif step == ExecutionStep.EXECUTE_ITERATION: + await self._step_execute_iteration(execution, iteration) + elif step == ExecutionStep.REVIEW_CHANGES: + await self._step_review_changes(execution, iteration) + elif step == ExecutionStep.COMMIT_CHANGES: + await self._step_commit_changes(execution, iteration) + elif step == ExecutionStep.HUMAN_INTERACTION: + await self._step_human_interaction(execution, iteration) + elif step == ExecutionStep.FINALIZE_TASK: + await self._step_finalize_task(execution, iteration) + + iteration.completed_at = datetime.now() + iteration.success = True + + except Exception as e: + iteration.completed_at = datetime.now() + iteration.success = False + iteration.error_message = str(e) + raise + + finally: + # Store iteration in database + await self._store_task_iteration(execution.task_id, iteration) + + async def _step_analyze_task( + self, execution: ExecutionContext, iteration: TaskIteration + ): + """Analyze the task and prepare execution plan""" + execution.status = TaskStatus.ANALYZING + await self._update_task_status(execution.task_id, TaskStatus.ANALYZING) + + # Analyze task requirements + task_description = execution.task_data.get("description", "") + task_title = execution.task_data.get("title", "") + + iteration.input_data = { + "task_title": task_title, + "task_description": task_description, + "agent_type": execution.agent_data.get("type"), + "agent_role": execution.agent_data.get("role"), + } + + # Enhance context with organizational knowledge + enhanced_context = None + if self.context_enhancer: + try: + enhanced_context = await self.context_enhancer.enhance_agent_context( + agent_id=execution.agent_id, + task_data=execution.task_data, + base_context=iteration.input_data, + ) + logger.info( + f"Enhanced context for task {execution.task_id}: " + f"{len(enhanced_context.organizational_knowledge)} org + " + f"{len(enhanced_context.team_knowledge)} team + " + f"{len(enhanced_context.similar_task_insights)} similar task insights" + ) + except Exception as e: + logger.error( + f"Error enhancing context for task {execution.task_id}: {e}" + ) + + # Simple analysis for now - in a full implementation this would use AI + iteration.output_data = { + "analysis_complete": True, + "requires_sandbox": execution.agent_data.get("type") == "developer", + "requires_git": bool( + execution.agent_data.get("repository_settings", {}).get( + "repository_url" + ) + ), + "enhanced_context": enhanced_context, + "estimated_complexity": "medium", + "estimated_iterations": 5, + } + + logger.info(f"Task analysis complete for {execution.task_id}") + + async def _step_setup_sandbox( + self, execution: ExecutionContext, iteration: TaskIteration + ): + """Set up sandbox environment for the agent""" + execution.status = TaskStatus.SETTING_UP + await self._update_task_status(execution.task_id, TaskStatus.SETTING_UP) + + agent_template = execution.agent_data.get("template_id", "python_developer") + repository_settings = execution.agent_data.get("repository_settings", {}) + sandbox_settings = execution.agent_data.get("sandbox_settings", {}) + + # Create sandbox + sandbox = await self.sandbox_manager.create_sandbox( + agent_id=execution.agent_id, + task_id=execution.task_id, + agent_template=agent_template, + repository_settings=repository_settings, + custom_settings=sandbox_settings, + ) + + execution.sandbox = sandbox + + iteration.input_data = { + "agent_template": agent_template, + "repository_settings": repository_settings, + "sandbox_settings": sandbox_settings, + } + + iteration.output_data = { + "sandbox_id": sandbox.sandbox_id, + "workspace_path": sandbox.workspace_path, + "container_id": sandbox.container_id, + } + + logger.info( + f"Sandbox setup complete for {execution.task_id}: {sandbox.sandbox_id}" + ) + + async def _step_setup_git( + self, execution: ExecutionContext, iteration: TaskIteration + ): + """Set up Git workflow for the task""" + repository_settings = execution.agent_data.get("repository_settings", {}) + + if not repository_settings.get("repository_url"): + # Skip Git setup if no repository + iteration.output_data = { + "git_setup": "skipped", + "reason": "no_repository_configured", + } + return + + # Create Git workflow manager + git_manager = GitWorkflowManager( + agent_id=execution.agent_id, + task_id=execution.task_id, + repo_settings=repository_settings, + ) + + # Setup workspace + feature_branch = await git_manager.setup_workspace() + + execution.git_manager = git_manager + + # Create enhanced Claude wrapper with Git context and conversation tracking + execution.claude_wrapper = ClaudeCodeWrapper( + workspace_path=git_manager.workspace_path, + git_manager=git_manager, + agent_id=execution.agent_id, + task_id=execution.task_id, + conversation_manager=self.conversation_manager, + ) + + # Initialize File Operations Engine + file_ops_engine = FileOperationsEngine(git_manager.workspace_path) + execution.file_operations_engine = file_ops_engine + self.file_operations_engines[execution.task_id] = file_ops_engine + + # Initialize Claude SDK Manager + claude_sdk_manager = ClaudeSDKManager( + file_operations_engine=file_ops_engine, + conversation_manager=self.conversation_manager, + ) + execution.claude_sdk_manager = claude_sdk_manager + self.claude_sdk_managers[execution.task_id] = claude_sdk_manager + + # Start conversation session + await execution.claude_wrapper.start_conversation_session( + execution.sandbox.sandbox_id + ) + + iteration.input_data = { + "repository_url": repository_settings.get("repository_url"), + "default_branch": repository_settings.get("default_branch", "main"), + } + + iteration.output_data = { + "git_setup": "complete", + "feature_branch": feature_branch, + "workspace_path": git_manager.workspace_path, + } + + logger.info(f"Git setup complete for {execution.task_id}: {feature_branch}") + + async def _step_execute_iteration( + self, execution: ExecutionContext, iteration: TaskIteration + ): + """Execute a development iteration using Claude SDK Manager""" + execution.status = TaskStatus.EXECUTING + await self._update_task_status(execution.task_id, TaskStatus.EXECUTING) + + task_description = execution.task_data.get("description", "") + task_title = execution.task_data.get("title", "") + + iteration.input_data = { + "task_description": task_description, + "task_title": task_title, + "iteration_number": iteration.iteration_number, + "workspace_path": ( + execution.sandbox.workspace_path if execution.sandbox else None + ), + } + + try: + # Start Claude SDK session if not already running + if not execution.claude_session_id and execution.claude_sdk_manager: + # Build comprehensive task context + context_info = "" + if execution.git_manager: + context_info += f"\nRepository: {execution.git_manager.repo_url}" + context_info += f"\nBranch: {execution.git_manager.feature_branch}" + if iteration.iteration_number > 1: + context_info += f"\nIteration: {iteration.iteration_number} of ongoing development" + + # Construct task prompt + task_prompt = f""" +Task: {task_title} + +Description: {task_description} + +Context: {context_info} + +Please analyze the codebase, understand the requirements, and implement the necessary changes. +Work incrementally and ask for clarification if needed. +""" + + # Start Claude SDK session + execution.claude_session_id = ( + await execution.claude_sdk_manager.start_session( + task_id=execution.task_id, + agent_id=execution.agent_id, + workspace_path=( + execution.git_manager.workspace_path + if execution.git_manager + else execution.sandbox.workspace_path + ), + task_description=task_prompt, + additional_context=context_info, + ) + ) + + logger.info( + f"Started Claude SDK session: {execution.claude_session_id}" + ) + + # Register interaction callback for human-in-the-loop + if execution.claude_sdk_manager and execution.claude_session_id: + execution.claude_sdk_manager.register_interaction_callback( + execution.claude_session_id, + lambda session, interaction: self._handle_claude_interaction( + execution, session, interaction + ), + ) + + # Monitor session status + session_status = ( + await execution.claude_sdk_manager.get_session_status( + execution.claude_session_id + ) + if execution.claude_session_id + else None + ) + + if session_status: + current_interaction = session_status.get("current_interaction") + if current_interaction: + # Claude is waiting for human input + interaction_type = current_interaction.get("type") + if interaction_type in ["user_input", "confirmation"]: + iteration.human_question = current_interaction.get("prompt") + execution.status = TaskStatus.WAITING_FOR_HUMAN + await self._update_task_status( + execution.task_id, TaskStatus.WAITING_FOR_HUMAN + ) + + iteration.output_data = { + "claude_session_status": session_status.get("state"), + "human_interaction_required": True, + "interaction_type": interaction_type, + "human_question_asked": True, + } + + logger.info( + f"Claude SDK requesting human input for task {execution.task_id}" + ) + return + + elif interaction_type == "file_approval": + # File operations pending approval + batch_id = current_interaction.get("metadata", {}).get( + "batch_id" + ) + if batch_id and execution.file_operations_engine: + # Get diff preview for human review + diffs = await execution.file_operations_engine.get_file_diff_preview( + batch_id + ) + + iteration.human_question = f"""Claude wants to make the following file changes: + +{current_interaction.get('prompt')} + +File changes preview: +{self._format_diffs_for_human(diffs)} + +Approve these changes? (yes/no)""" + + execution.status = TaskStatus.WAITING_FOR_HUMAN + await self._update_task_status( + execution.task_id, TaskStatus.WAITING_FOR_HUMAN + ) + + iteration.output_data = { + "claude_session_status": session_status.get("state"), + "file_approval_required": True, + "batch_id": batch_id, + "file_changes_preview": diffs, + "human_question_asked": True, + } + + logger.info( + f"Claude SDK requesting file approval for task {execution.task_id}" + ) + return + + # No interaction needed - continue execution + iteration.output_data = { + "claude_session_status": session_status.get("state"), + "session_active": True, + "iteration_completed": True, + "workspace_path": session_status.get("workspace_path"), + } + + # Check if session completed + if session_status.get("state") in ["completed", "terminated"]: + iteration.output_data["development_complete"] = True + + else: + # No session - this shouldn't happen but handle gracefully + iteration.output_data = { + "error": "No Claude SDK session available", + "development_complete": False, + } + + except Exception as e: + logger.error(f"Error in Claude SDK iteration: {e}") + iteration.output_data = {"error": str(e), "development_complete": False} + raise + + logger.info( + f"Development iteration {iteration.iteration_number} processed for {execution.task_id}" + ) + + async def _step_review_changes( + self, execution: ExecutionContext, iteration: TaskIteration + ): + """Review the changes made in the iteration""" + execution.status = TaskStatus.REVIEWING + await self._update_task_status(execution.task_id, TaskStatus.REVIEWING) + + # Review changes - in full implementation this would: + # 1. Run linting and type checking + # 2. Run tests + # 3. Check code quality + # 4. Validate against requirements + + iteration.output_data = { + "review_passed": True, + "issues_found": [], + "tests_passed": True, + "code_quality_score": 85, + } + + logger.info(f"Code review complete for {execution.task_id}") + + async def _step_commit_changes( + self, execution: ExecutionContext, iteration: TaskIteration + ): + """Commit changes to Git""" + execution.status = TaskStatus.COMMITTING + await self._update_task_status(execution.task_id, TaskStatus.COMMITTING) + + if not execution.git_manager: + iteration.output_data = {"commit": "skipped", "reason": "no_git_manager"} + return + + # Commit changes + commit_message = f"Iteration {iteration.iteration_number}: {execution.task_data.get('title', 'Task update')}" + commit_hash = await execution.git_manager.commit_changes( + message=commit_message, iteration_number=iteration.iteration_number + ) + + iteration.output_data = { + "commit_hash": commit_hash, + "commit_message": commit_message, + "branch": execution.git_manager.feature_branch, + } + + logger.info(f"Changes committed for {execution.task_id}: {commit_hash}") + + async def _step_human_interaction( + self, execution: ExecutionContext, iteration: TaskIteration + ): + """Handle human interaction step""" + # This step waits for human response - the actual waiting is handled + # by the status being WAITING_FOR_HUMAN + iteration.output_data = { + "human_interaction": "waiting_for_response", + "question": iteration.human_question, + } + + async def _step_finalize_task( + self, execution: ExecutionContext, iteration: TaskIteration + ): + """Finalize the task execution""" + execution.status = TaskStatus.COMPLETED + execution.completed_at = datetime.now() + + # Create pull request if Git is configured + pr_url = None + if execution.git_manager: + try: + pr = await execution.git_manager.create_pull_request( + title=f"🤖 {execution.task_data.get('title', 'Task completion')}", + description=f"Autonomous completion of task: {execution.task_data.get('description', '')}", + ) + pr_url = pr.url if pr else None + except Exception as e: + logger.error(f"Failed to create PR for {execution.task_id}: {e}") + + # Prepare final result + execution.result = { + "status": "completed", + "iterations": len(execution.iterations), + "started_at": execution.started_at.isoformat(), + "completed_at": execution.completed_at.isoformat(), + "pull_request_url": pr_url, + "sandbox_id": execution.sandbox.sandbox_id if execution.sandbox else None, + "git_branch": ( + execution.git_manager.feature_branch if execution.git_manager else None + ), + } + + # Update database + await self._update_task_status( + execution.task_id, TaskStatus.COMPLETED, result=execution.result + ) + + iteration.output_data = execution.result + + # End conversation session + if execution.claude_wrapper: + try: + await execution.claude_wrapper.end_conversation_session() + except Exception as e: + logger.error(f"Error ending conversation session: {e}") + + # Store final performance metrics + await self._store_completion_metrics(execution) + + # Extract knowledge from completed task + if self.knowledge_extractor: + try: + extracted_knowledge_ids = ( + await self.knowledge_extractor.extract_knowledge_from_task( + task_id=execution.task_id, + agent_id=execution.agent_id, + execution_result=execution.result, + ) + ) + if extracted_knowledge_ids: + logger.info( + f"Extracted {len(extracted_knowledge_ids)} knowledge items from task {execution.task_id}" + ) + except Exception as e: + logger.error( + f"Error extracting knowledge from task {execution.task_id}: {e}" + ) + + logger.info(f"✅ Task execution completed: {execution.task_id}") + + # End Claude SDK session + if execution.claude_sdk_manager and execution.claude_session_id: + try: + await execution.claude_sdk_manager.terminate_session( + execution.claude_session_id + ) + except Exception as e: + logger.error(f"Error terminating Claude SDK session: {e}") + + # Schedule cleanup + asyncio.create_task(self._cleanup_execution(execution.task_id)) + + def _is_task_complete(self, execution: ExecutionContext) -> bool: + """Check if the task is complete""" + # Simple completion check - in full implementation this would be more sophisticated + return execution.current_iteration >= 3 + + async def _handle_execution_error(self, task_id: str, error_message: str): + """Handle execution error""" + execution = self.active_executions.get(task_id) + if execution: + execution.status = TaskStatus.FAILED + execution.completed_at = datetime.now() + execution.error = error_message + + await self._update_task_status(task_id, TaskStatus.FAILED, error=error_message) + + # Schedule cleanup + asyncio.create_task(self._cleanup_execution(task_id)) + + logger.error(f"Task execution failed: {task_id} - {error_message}") + + async def _cleanup_execution(self, task_id: str): + """Clean up execution resources""" + execution = self.active_executions.pop(task_id, None) + if not execution: + return + + # Cleanup Claude SDK session + if execution.claude_sdk_manager and execution.claude_session_id: + try: + await execution.claude_sdk_manager.terminate_session( + execution.claude_session_id + ) + except Exception as e: + logger.error(f"Error terminating Claude SDK session for {task_id}: {e}") + + # Remove engines from tracking + self.file_operations_engines.pop(task_id, None) + self.claude_sdk_managers.pop(task_id, None) + + # Cleanup sandbox + if execution.sandbox: + try: + await self.sandbox_manager.destroy_sandbox(execution.sandbox.sandbox_id) + except Exception as e: + logger.error(f"Error destroying sandbox for {task_id}: {e}") + + # Cleanup Git workspace (optional - might want to keep for review) + if execution.git_manager and execution.status != TaskStatus.COMPLETED: + try: + await execution.git_manager.cleanup_workspace() + except Exception as e: + logger.error(f"Error cleaning up Git workspace for {task_id}: {e}") + + logger.info(f"Execution cleanup complete: {task_id}") + + async def _monitoring_worker(self): + """Monitor execution health and timeouts""" + while self.running: + try: + current_time = datetime.now() + + for task_id, execution in list(self.active_executions.items()): + # Check for timeouts + if execution.status == TaskStatus.WAITING_FOR_HUMAN: + # Check human response timeout + last_iteration = ( + execution.iterations[-1] if execution.iterations else None + ) + if last_iteration and last_iteration.started_at: + time_waiting = current_time - last_iteration.started_at + if time_waiting > timedelta( + seconds=self.human_response_timeout + ): + await self._handle_execution_error( + task_id, + "Human response timeout - no response received within 24 hours", + ) + else: + # Check general execution timeout + execution_time = current_time - execution.started_at + if execution_time > timedelta( + seconds=self.iteration_timeout * self.max_iterations + ): + await self._handle_execution_error( + task_id, + f"Execution timeout - exceeded maximum time limit", + ) + + # Check iteration limits + if execution.current_iteration > self.max_iterations: + await self._handle_execution_error( + task_id, + f"Maximum iterations exceeded ({self.max_iterations})", + ) + + await asyncio.sleep(60) # Check every minute + + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"Error in monitoring worker: {e}") + await asyncio.sleep(60) + + async def _cleanup_worker(self): + """Clean up completed executions periodically""" + while self.running: + try: + current_time = datetime.now() + cutoff_time = current_time - timedelta( + hours=1 + ) # Keep completed tasks for 1 hour + + completed_tasks = [ + task_id + for task_id, execution in self.active_executions.items() + if execution.status + in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED] + and execution.completed_at + and execution.completed_at < cutoff_time + ] + + for task_id in completed_tasks: + await self._cleanup_execution(task_id) + + await asyncio.sleep(3600) # Run every hour + + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"Error in cleanup worker: {e}") + await asyncio.sleep(3600) + + # Database operations + + async def _get_task_data(self, task_id: str) -> Optional[Dict[str, Any]]: + """Get task data from database""" + async with get_db_connection() as conn: + row = await conn.fetchrow("SELECT * FROM tasks WHERE id = $1", task_id) + return dict(row) if row else None + + async def _get_agent_data(self, agent_id: str) -> Optional[Dict[str, Any]]: + """Get agent data from database""" + return await DatabaseManager.get_agent(agent_id) + + async def _update_task_status( + self, + task_id: str, + status: TaskStatus, + result: Optional[Dict[str, Any]] = None, + error: Optional[str] = None, + ): + """Update task status in database""" + await DatabaseManager.update_task_status( + task_id=task_id, status=status.value, result=result + ) + + async def _store_task_iteration(self, task_id: str, iteration: TaskIteration): + """Store task iteration in database""" + async with get_db_connection() as conn: + await conn.execute( + """ + INSERT INTO task_iterations ( + id, task_id, iteration_number, step, started_at, completed_at, + input_data, output_data, success, error_message, + human_question, human_response + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) + """, + str(uuid.uuid4()), + task_id, + iteration.iteration_number, + iteration.step.value, + iteration.started_at, + iteration.completed_at, + json.dumps(iteration.input_data), + json.dumps(iteration.output_data) if iteration.output_data else None, + iteration.success, + iteration.error_message, + iteration.human_question, + iteration.human_response, + ) + + async def _get_task_iterations_from_db(self, task_id: str) -> List[TaskIteration]: + """Get task iterations from database""" + async with get_db_connection() as conn: + rows = await conn.fetch( + """ + SELECT * FROM task_iterations + WHERE task_id = $1 + ORDER BY iteration_number + """, + task_id, + ) + + iterations = [] + for row in rows: + iterations.append( + TaskIteration( + iteration_number=row["iteration_number"], + step=ExecutionStep(row["step"]), + started_at=row["started_at"], + completed_at=row["completed_at"], + input_data=( + json.loads(row["input_data"]) if row["input_data"] else {} + ), + output_data=( + json.loads(row["output_data"]) + if row["output_data"] + else None + ), + success=row["success"], + error_message=row["error_message"], + human_question=row["human_question"], + human_response=row["human_response"], + ) + ) + + return iterations + + async def _store_completion_metrics(self, execution: ExecutionContext): + """Store performance metrics for completed task""" + try: + if not execution.completed_at or not execution.started_at: + return + + # Calculate execution time + execution_time = ( + execution.completed_at - execution.started_at + ).total_seconds() / 60 # minutes + + # Store metrics + await self.conversation_manager.store_performance_metric( + agent_id=execution.agent_id, + task_id=execution.task_id, + metric_type="execution_time_minutes", + metric_value=execution_time, + metric_unit="minutes", + ) + + await self.conversation_manager.store_performance_metric( + agent_id=execution.agent_id, + task_id=execution.task_id, + metric_type="iterations_to_completion", + metric_value=float(execution.current_iteration), + metric_unit="iterations", + ) + + # Store success/failure metric + success_value = 1.0 if execution.status == TaskStatus.COMPLETED else 0.0 + await self.conversation_manager.store_performance_metric( + agent_id=execution.agent_id, + task_id=execution.task_id, + metric_type="task_success_rate", + metric_value=success_value, + metric_unit="boolean", + ) + + except Exception as e: + logger.error(f"Error storing completion metrics: {e}") + + async def _handle_claude_interaction( + self, execution: ExecutionContext, session: ClaudeSDKSession, interaction + ): + """Handle interaction from Claude SDK""" + logger.info( + f"Claude interaction for task {execution.task_id}: {interaction.interaction_type}" + ) + + # This will be processed in the next iteration of _step_execute_iteration + # The interaction handling is done there to maintain the execution flow + + def _format_diffs_for_human(self, diffs: Dict[str, str]) -> str: + """Format file diffs for human review""" + if not diffs: + return "No file changes detected." + + formatted = [] + for file_path, diff in diffs.items(): + formatted.append(f"\n--- {file_path} ---") + formatted.append(diff[:1000] + "..." if len(diff) > 1000 else diff) + + return "\n".join(formatted) + + async def handle_human_response(self, task_id: str, response: str) -> bool: + """Enhanced human response handler that integrates with Claude SDK""" + + execution = self.active_executions.get(task_id) + if not execution or execution.status != TaskStatus.WAITING_FOR_HUMAN: + return False + + # Find the current iteration waiting for human response + current_iteration = None + for iteration in reversed(execution.iterations): + if iteration.human_question and not iteration.human_response: + current_iteration = iteration + break + + if current_iteration: + current_iteration.human_response = response + execution.status = TaskStatus.EXECUTING + + # Handle different types of responses + output_data = current_iteration.output_data or {} + + if output_data.get("file_approval_required"): + # Handle file approval + batch_id = output_data.get("batch_id") + approved = response.lower().strip() in [ + "yes", + "y", + "approve", + "approved", + "true", + ] + + if ( + batch_id + and execution.claude_sdk_manager + and execution.claude_session_id + ): + success = ( + await execution.claude_sdk_manager.approve_file_operations( + execution.claude_session_id, batch_id, approved + ) + ) + + if success: + logger.info( + f"File operations {'approved' if approved else 'rejected'} for task {task_id}" + ) + else: + logger.error( + f"Failed to process file approval for task {task_id}" + ) + + else: + # Handle general user input + if execution.claude_sdk_manager and execution.claude_session_id: + success = await execution.claude_sdk_manager.send_input( + execution.claude_session_id, response + ) + + if success: + logger.info( + f"Human response sent to Claude SDK for task {task_id}" + ) + else: + logger.error( + f"Failed to send human response to Claude SDK for task {task_id}" + ) + + # Update database + await self._store_task_iteration(task_id, current_iteration) + await self._update_task_status(task_id, TaskStatus.EXECUTING) + + logger.info(f"Human response processed for task {task_id}") + return True + + return False diff --git a/services/orchestrator/task_knowledge_extractor.py b/services/orchestrator/task_knowledge_extractor.py index 80501ba..409755d 100644 --- a/services/orchestrator/task_knowledge_extractor.py +++ b/services/orchestrator/task_knowledge_extractor.py @@ -1,879 +1,879 @@ -""" -Task Knowledge Extractor for FuzeAgent - -This module extracts valuable knowledge from completed tasks and feeds it -into the hierarchical knowledge management system. It analyzes task outcomes, -code changes, conversation patterns, and performance metrics to create -reusable organizational knowledge. -""" - -import asyncio -import json -import logging -import re -from dataclasses import dataclass -from datetime import datetime -from typing import Any, Dict, List, Optional, Tuple - -import asyncpg -from sentence_transformers import SentenceTransformer - -from .knowledge_propagation_engine import KnowledgePropagationEngine, PropagationTrigger -from .organization_rag_manager import ( - ContentType, - KnowledgeCategory, - OrganizationRAGManager, - SourceType, -) -from .team_knowledge_manager import TeamKnowledgeManager - -logger = logging.getLogger(__name__) - - -@dataclass -class TaskKnowledgeExtract: - """Represents extracted knowledge from a task""" - - title: str - content: str - content_type: ContentType - category: KnowledgeCategory - confidence_score: float - tags: List[str] - metadata: Dict[str, Any] - success_indicators: List[str] - failure_patterns: List[str] - - -@dataclass -class ExtractionContext: - """Context for knowledge extraction""" - - task_id: str - agent_id: str - team_id: str - organization_id: str - task_data: Dict[str, Any] - execution_result: Dict[str, Any] - conversation_history: List[Dict[str, Any]] - code_changes: List[Dict[str, Any]] - performance_metrics: Dict[str, Any] - iteration_count: int - total_duration_minutes: float - success: bool - - -class TaskKnowledgeExtractor: - """ - Extracts knowledge from completed tasks and integrates it - into the hierarchical knowledge management system. - """ - - def __init__( - self, - database_url: str, - org_rag_manager: OrganizationRAGManager, - team_knowledge_manager: TeamKnowledgeManager, - propagation_engine: KnowledgePropagationEngine, - ): - self.database_url = database_url - self.org_rag_manager = org_rag_manager - self.team_knowledge_manager = team_knowledge_manager - self.propagation_engine = propagation_engine - self.pool: Optional[asyncpg.Pool] = None - - # Initialize text analysis model - self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") - - # Extraction patterns and rules - self.code_patterns = self._initialize_code_patterns() - self.success_patterns = self._initialize_success_patterns() - self.failure_patterns = self._initialize_failure_patterns() - - # Configuration - self.min_extraction_confidence = 0.4 - self.min_task_duration_minutes = 5 # Don't extract from very short tasks - self.max_content_length = 5000 - - # Statistics - self.extractions_performed = 0 - self.knowledge_items_created = 0 - self.propagations_triggered = 0 - - async def initialize(self): - """Initialize the knowledge extractor""" - logger.info("Initializing TaskKnowledgeExtractor") - - try: - self.pool = await asyncpg.create_pool( - self.database_url, min_size=1, max_size=5, command_timeout=60 - ) - - logger.info("TaskKnowledgeExtractor initialized successfully") - - except Exception as e: - logger.error(f"Failed to initialize TaskKnowledgeExtractor: {e}") - raise - - async def close(self): - """Close database connections""" - if self.pool: - await self.pool.close() - logger.info("TaskKnowledgeExtractor closed") - - async def extract_knowledge_from_task( - self, task_id: str, agent_id: str, execution_result: Dict[str, Any] - ) -> List[str]: - """Extract knowledge from a completed task and store it""" - - try: - # Build extraction context - context = await self._build_extraction_context( - task_id, agent_id, execution_result - ) - - if not context: - logger.warning(f"Could not build extraction context for task {task_id}") - return [] - - # Skip extraction for very short or trivial tasks - if context.total_duration_minutes < self.min_task_duration_minutes: - logger.debug( - f"Skipping extraction for short task {task_id} ({context.total_duration_minutes:.1f}m)" - ) - return [] - - # Extract knowledge items - knowledge_extracts = await self._extract_knowledge_items(context) - - if not knowledge_extracts: - logger.debug(f"No knowledge extracted from task {task_id}") - return [] - - # Store extracted knowledge - stored_knowledge_ids = [] - for extract in knowledge_extracts: - if extract.confidence_score >= self.min_extraction_confidence: - knowledge_id = await self._store_knowledge_extract(context, extract) - if knowledge_id: - stored_knowledge_ids.append(knowledge_id) - - # Trigger knowledge propagation if we have valuable knowledge - if stored_knowledge_ids: - propagation_ids = ( - await self.propagation_engine.trigger_agent_to_team_propagation( - agent_id=context.agent_id, - task_id=context.task_id, - task_outcome={ - "success": context.success, - "task_type": context.task_data.get("task_type", "unknown"), - "complexity": self._assess_task_complexity(context), - "duration_minutes": context.total_duration_minutes, - "knowledge_extracted": len(stored_knowledge_ids), - "iteration_count": context.iteration_count, - }, - ) - ) - - self.propagations_triggered += len(propagation_ids) - logger.info( - f"Triggered {len(propagation_ids)} knowledge propagations for task {task_id}" - ) - - self.extractions_performed += 1 - self.knowledge_items_created += len(stored_knowledge_ids) - - logger.info( - f"Extracted {len(stored_knowledge_ids)} knowledge items from task {task_id}" - ) - return stored_knowledge_ids - - except Exception as e: - logger.error(f"Error extracting knowledge from task {task_id}: {e}") - return [] - - async def _build_extraction_context( - self, task_id: str, agent_id: str, execution_result: Dict[str, Any] - ) -> Optional[ExtractionContext]: - """Build context for knowledge extraction""" - - async with self.pool.acquire() as conn: - # Get basic task information - task_data = await conn.fetchrow( - """ - SELECT t.*, a.team_id, te.organization_id - FROM tasks t - JOIN agents a ON t.agent_id = a.id - JOIN teams te ON a.team_id = te.id - WHERE t.id = $1 - """, - task_id, - ) - - if not task_data: - return None - - # Get conversation history - conversation_history = await conn.fetch( - """ - SELECT message_type, content, metadata, created_at - FROM claude_conversations - WHERE task_id = $1 - ORDER BY created_at ASC - """, - task_id, - ) - - # Get code generations - code_changes = await conn.fetch( - """ - SELECT file_path, file_type, language, content, test_results, quality_metrics - FROM code_generations - WHERE task_id = $1 - ORDER BY generated_at ASC - """, - task_id, - ) - - # Get performance metrics - performance_metrics = await conn.fetch( - """ - SELECT metric_type, metric_value, metric_unit, context - FROM agent_performance_metrics - WHERE task_id = $1 - """, - task_id, - ) - - # Calculate duration - started_at = task_data["started_at"] - completed_at = execution_result.get("completed_at") - if completed_at: - if isinstance(completed_at, str): - completed_at = datetime.fromisoformat( - completed_at.replace("Z", "+00:00") - ) - duration = (completed_at - started_at).total_seconds() / 60.0 - else: - duration = 0.0 - - return ExtractionContext( - task_id=str(task_data["id"]), - agent_id=str(task_data["agent_id"]), - team_id=str(task_data["team_id"]), - organization_id=str(task_data["organization_id"]), - task_data=dict(task_data), - execution_result=execution_result, - conversation_history=[dict(conv) for conv in conversation_history], - code_changes=[dict(code) for code in code_changes], - performance_metrics={ - pm["metric_type"]: pm for pm in performance_metrics - }, - iteration_count=execution_result.get("iterations", 0), - total_duration_minutes=duration, - success=execution_result.get("status") == "completed", - ) - - async def _extract_knowledge_items( - self, context: ExtractionContext - ) -> List[TaskKnowledgeExtract]: - """Extract specific knowledge items from the task context""" - - knowledge_extracts = [] - - # Extract different types of knowledge - knowledge_extracts.extend(await self._extract_code_patterns(context)) - knowledge_extracts.extend(await self._extract_problem_solutions(context)) - knowledge_extracts.extend(await self._extract_debugging_insights(context)) - knowledge_extracts.extend(await self._extract_process_knowledge(context)) - knowledge_extracts.extend(await self._extract_error_patterns(context)) - knowledge_extracts.extend(await self._extract_optimization_insights(context)) - - return knowledge_extracts - - async def _extract_code_patterns( - self, context: ExtractionContext - ) -> List[TaskKnowledgeExtract]: - """Extract reusable code patterns and best practices""" - - extracts = [] - - for code_change in context.code_changes: - if code_change["file_type"] == "implementation": - content = code_change["content"] - language = code_change.get("language", "unknown") - - # Look for reusable patterns - patterns_found = [] - for pattern_name, pattern_info in self.code_patterns.items(): - if any( - keyword in content.lower() - for keyword in pattern_info["keywords"] - ): - patterns_found.append(pattern_name) - - if patterns_found and len(content) > 100: # Substantial code - # Create knowledge extract - title = f"Code Pattern: {', '.join(patterns_found)} ({language})" - extract_content = self._create_code_pattern_content( - content, patterns_found, context - ) - - confidence = self._calculate_code_pattern_confidence( - content, patterns_found, context - ) - - if confidence >= self.min_extraction_confidence: - extract = TaskKnowledgeExtract( - title=title, - content=extract_content, - content_type=ContentType.CODE, - category=KnowledgeCategory.DEVELOPMENT, - confidence_score=confidence, - tags=["code_pattern", language, *patterns_found], - metadata={ - "language": language, - "file_path": code_change["file_path"], - "patterns": patterns_found, - "task_success": context.success, - "lines_of_code": len(content.split("\n")), - }, - success_indicators=self._extract_success_indicators( - context - ), - failure_patterns=[], - ) - - extracts.append(extract) - - return extracts - - async def _extract_problem_solutions( - self, context: ExtractionContext - ) -> List[TaskKnowledgeExtract]: - """Extract problem-solution pairs from the task""" - - extracts = [] - - # Analyze conversation for problem descriptions and solutions - problem_solution_pairs = self._identify_problem_solution_pairs( - context.conversation_history - ) - - for problem, solution in problem_solution_pairs: - if len(problem) > 50 and len(solution) > 50: # Substantial content - title = f"Solution: {problem[:50]}..." - content = f"**Problem:**\n{problem}\n\n**Solution:**\n{solution}" - - # Determine category based on content - category = self._categorize_problem_solution(problem, solution) - - confidence = self._calculate_solution_confidence( - problem, solution, context - ) - - if confidence >= self.min_extraction_confidence: - extract = TaskKnowledgeExtract( - title=title, - content=content[: self.max_content_length], - content_type=ContentType.PROCEDURE, - category=category, - confidence_score=confidence, - tags=["problem_solution", "troubleshooting"], - metadata={ - "problem_type": self._classify_problem_type(problem), - "solution_type": self._classify_solution_type(solution), - "task_success": context.success, - }, - success_indicators=self._extract_success_indicators(context), - failure_patterns=[], - ) - - extracts.append(extract) - - return extracts - - async def _extract_debugging_insights( - self, context: ExtractionContext - ) -> List[TaskKnowledgeExtract]: - """Extract debugging approaches and insights""" - - extracts = [] - - # Look for error messages and resolution patterns - debugging_sessions = self._identify_debugging_sessions( - context.conversation_history - ) - - for session in debugging_sessions: - if session["resolution"] and context.success: - title = f"Debugging: {session['error_type']}" - content = self._create_debugging_content(session) - - confidence = self._calculate_debugging_confidence(session, context) - - if confidence >= self.min_extraction_confidence: - extract = TaskKnowledgeExtract( - title=title, - content=content, - content_type=ContentType.PROCEDURE, - category=KnowledgeCategory.TROUBLESHOOTING, - confidence_score=confidence, - tags=["debugging", session["error_type"], "troubleshooting"], - metadata={ - "error_type": session["error_type"], - "resolution_method": session["resolution_method"], - "tools_used": session.get("tools_used", []), - }, - success_indicators=self._extract_success_indicators(context), - failure_patterns=session.get("failure_patterns", []), - ) - - extracts.append(extract) - - return extracts - - async def _extract_process_knowledge( - self, context: ExtractionContext - ) -> List[TaskKnowledgeExtract]: - """Extract process and workflow knowledge""" - - extracts = [] - - if context.iteration_count > 1: # Multi-iteration tasks have process insights - title = f"Process: {context.task_data.get('task_type', 'Task')} Workflow" - - process_content = self._create_process_content(context) - confidence = self._calculate_process_confidence(context) - - if confidence >= self.min_extraction_confidence: - extract = TaskKnowledgeExtract( - title=title, - content=process_content, - content_type=ContentType.PROCEDURE, - category=KnowledgeCategory.PROCESS, - confidence_score=confidence, - tags=[ - "process", - "workflow", - context.task_data.get("task_type", "general"), - ], - metadata={ - "iterations_used": context.iteration_count, - "duration_minutes": context.total_duration_minutes, - "success_rate": 1.0 if context.success else 0.0, - "complexity": self._assess_task_complexity(context), - }, - success_indicators=self._extract_success_indicators(context), - failure_patterns=[], - ) - - extracts.append(extract) - - return extracts - - async def _extract_error_patterns( - self, context: ExtractionContext - ) -> List[TaskKnowledgeExtract]: - """Extract error patterns and avoidance strategies""" - - extracts = [] - - # Look for error patterns in failed tasks or recovered errors - error_patterns = self._identify_error_patterns(context.conversation_history) - - for pattern in error_patterns: - if pattern["frequency"] >= 2 or pattern["severity"] == "high": - title = f"Error Pattern: {pattern['error_type']}" - content = self._create_error_pattern_content(pattern, context) - - confidence = self._calculate_error_pattern_confidence(pattern, context) - - if confidence >= self.min_extraction_confidence: - extract = TaskKnowledgeExtract( - title=title, - content=content, - content_type=ContentType.DOCUMENTATION, - category=KnowledgeCategory.TROUBLESHOOTING, - confidence_score=confidence, - tags=["error_pattern", pattern["error_type"], "prevention"], - metadata={ - "error_type": pattern["error_type"], - "frequency": pattern["frequency"], - "severity": pattern["severity"], - "prevention_strategies": pattern.get("prevention", []), - }, - success_indicators=[], - failure_patterns=pattern.get("indicators", []), - ) - - extracts.append(extract) - - return extracts - - async def _extract_optimization_insights( - self, context: ExtractionContext - ) -> List[TaskKnowledgeExtract]: - """Extract performance optimization insights""" - - extracts = [] - - # Look for performance improvements in metrics - if "execution_time_minutes" in context.performance_metrics: - perf_data = context.performance_metrics["execution_time_minutes"] - if ( - perf_data["metric_value"] < 30 and context.success - ): # Efficient completion - title = "Performance Optimization: Efficient Task Execution" - content = self._create_optimization_content(context) - - confidence = self._calculate_optimization_confidence(context) - - if confidence >= self.min_extraction_confidence: - extract = TaskKnowledgeExtract( - title=title, - content=content, - content_type=ContentType.BEST_PRACTICE, - category=KnowledgeCategory.DEVELOPMENT, - confidence_score=confidence, - tags=["optimization", "performance", "efficiency"], - metadata={ - "execution_time": perf_data["metric_value"], - "iteration_efficiency": context.iteration_count - / context.total_duration_minutes, - "optimization_techniques": self._identify_optimization_techniques( - context - ), - }, - success_indicators=self._extract_success_indicators(context), - failure_patterns=[], - ) - - extracts.append(extract) - - return extracts - - async def _store_knowledge_extract( - self, context: ExtractionContext, extract: TaskKnowledgeExtract - ) -> Optional[str]: - """Store a knowledge extract in the appropriate knowledge base""" - - try: - # Store in organization knowledge base - knowledge_id = await self.org_rag_manager.add_knowledge( - organization_id=context.organization_id, - title=extract.title, - content=extract.content, - content_type=extract.content_type, - knowledge_category=extract.category, - source_type=SourceType.TASK_OUTCOME, - source_agent_id=context.agent_id, - source_team_id=context.team_id, - source_task_id=context.task_id, - relevance_score=extract.confidence_score, - quality_score=extract.confidence_score, - metadata={ - **extract.metadata, - "extraction_timestamp": datetime.now().isoformat(), - "extractor_version": "1.0", - "success_indicators": extract.success_indicators, - "failure_patterns": extract.failure_patterns, - }, - tags=extract.tags, - ) - - return knowledge_id - - except Exception as e: - logger.error(f"Error storing knowledge extract: {e}") - return None - - # Helper methods for pattern matching and analysis - def _initialize_code_patterns(self) -> Dict[str, Dict[str, Any]]: - """Initialize code pattern definitions""" - return { - "api_integration": { - "keywords": ["fetch", "axios", "request", "api", "endpoint", "rest"], - "confidence_boost": 0.2, - }, - "database_operations": { - "keywords": [ - "select", - "insert", - "update", - "delete", - "query", - "database", - "sql", - ], - "confidence_boost": 0.2, - }, - "authentication": { - "keywords": ["auth", "login", "token", "jwt", "session", "passport"], - "confidence_boost": 0.15, - }, - "error_handling": { - "keywords": ["try", "catch", "error", "exception", "throw"], - "confidence_boost": 0.1, - }, - "testing": { - "keywords": ["test", "spec", "describe", "it", "expect", "mock"], - "confidence_boost": 0.15, - }, - "optimization": { - "keywords": ["performance", "optimize", "cache", "memory", "speed"], - "confidence_boost": 0.2, - }, - } - - def _initialize_success_patterns(self) -> List[str]: - """Initialize success indicator patterns""" - return [ - r"test.*pass", - r"build.*success", - r"deploy.*complete", - r"fix.*issue", - r"resolve.*problem", - r"implement.*feature", - r"complete.*task", - ] - - def _initialize_failure_patterns(self) -> List[str]: - """Initialize failure indicator patterns""" - return [ - r"error.*occur", - r"fail.*to", - r"timeout.*exceed", - r"connection.*refuse", - r"not.*found", - r"access.*deni", - r"invalid.*request", - ] - - def _assess_task_complexity(self, context: ExtractionContext) -> str: - """Assess task complexity based on various factors""" - - complexity_score = 0 - - # Factor 1: Iteration count - if context.iteration_count > 10: - complexity_score += 3 - elif context.iteration_count > 5: - complexity_score += 2 - elif context.iteration_count > 2: - complexity_score += 1 - - # Factor 2: Duration - if context.total_duration_minutes > 180: # 3 hours - complexity_score += 3 - elif context.total_duration_minutes > 60: # 1 hour - complexity_score += 2 - elif context.total_duration_minutes > 30: - complexity_score += 1 - - # Factor 3: Code changes - if len(context.code_changes) > 10: - complexity_score += 2 - elif len(context.code_changes) > 5: - complexity_score += 1 - - # Factor 4: Conversation length - if len(context.conversation_history) > 50: - complexity_score += 2 - elif len(context.conversation_history) > 20: - complexity_score += 1 - - if complexity_score >= 6: - return "very_high" - elif complexity_score >= 4: - return "high" - elif complexity_score >= 2: - return "medium" - else: - return "low" - - def _extract_success_indicators(self, context: ExtractionContext) -> List[str]: - """Extract success indicators from the task execution""" - - indicators = [] - - # Look for success patterns in conversation - for conv in context.conversation_history: - content = conv.get("content", "").lower() - for pattern in self.success_patterns: - if re.search(pattern, content): - indicators.append(pattern) - - # Add task-specific indicators - if context.success: - indicators.append("task_completed_successfully") - - if context.execution_result.get("pull_request_url"): - indicators.append("pull_request_created") - - return list(set(indicators)) # Remove duplicates - - # Additional helper methods would be implemented here... - # (The file is getting quite long, so I'll implement key methods and indicate where others would go) - - def _identify_problem_solution_pairs( - self, conversation_history: List[Dict] - ) -> List[Tuple[str, str]]: - """Identify problem-solution pairs in conversation history""" - pairs = [] - # Implementation would analyze conversation flow to identify problems and their solutions - # This is a simplified placeholder - return pairs - - def _categorize_problem_solution( - self, problem: str, solution: str - ) -> KnowledgeCategory: - """Categorize a problem-solution pair""" - # Simple categorization based on keywords - combined_text = (problem + " " + solution).lower() - - if any(word in combined_text for word in ["test", "testing", "spec"]): - return KnowledgeCategory.TESTING - elif any(word in combined_text for word in ["deploy", "build", "ci", "cd"]): - return KnowledgeCategory.INFRASTRUCTURE - elif any(word in combined_text for word in ["security", "auth", "permission"]): - return KnowledgeCategory.SECURITY - elif any(word in combined_text for word in ["design", "ui", "ux", "interface"]): - return KnowledgeCategory.DESIGN - else: - return KnowledgeCategory.DEVELOPMENT - - def _calculate_code_pattern_confidence( - self, content: str, patterns: List[str], context: ExtractionContext - ) -> float: - """Calculate confidence score for code pattern extraction""" - base_confidence = 0.5 - - # Boost for successful task - if context.success: - base_confidence += 0.2 - - # Boost for multiple patterns - if len(patterns) > 1: - base_confidence += 0.1 - - # Boost for substantial code - if len(content) > 500: - base_confidence += 0.1 - - return min(1.0, base_confidence) - - def _calculate_solution_confidence( - self, problem: str, solution: str, context: ExtractionContext - ) -> float: - """Calculate confidence score for solution extraction""" - base_confidence = 0.4 - - if context.success: - base_confidence += 0.3 - - if len(solution) > 200: # Detailed solution - base_confidence += 0.1 - - return min(1.0, base_confidence) - - def _calculate_debugging_confidence( - self, session: Dict, context: ExtractionContext - ) -> float: - """Calculate confidence for debugging insights""" - base_confidence = 0.6 if context.success else 0.3 - - if session.get("resolution_method"): - base_confidence += 0.2 - - return min(1.0, base_confidence) - - def _calculate_process_confidence(self, context: ExtractionContext) -> float: - """Calculate confidence for process knowledge""" - if not context.success: - return 0.2 - - # Base confidence increases with iteration count (more process learning) - base_confidence = min(0.8, 0.3 + (context.iteration_count * 0.05)) - - return base_confidence - - def _calculate_error_pattern_confidence( - self, pattern: Dict, context: ExtractionContext - ) -> float: - """Calculate confidence for error pattern extraction""" - base_confidence = 0.4 - - if pattern["frequency"] > 2: - base_confidence += 0.2 - - if pattern["severity"] == "high": - base_confidence += 0.2 - - return min(1.0, base_confidence) - - def _calculate_optimization_confidence(self, context: ExtractionContext) -> float: - """Calculate confidence for optimization insights""" - if not context.success: - return 0.1 - - base_confidence = 0.5 - - # Boost for efficient execution - if context.total_duration_minutes < 30: - base_confidence += 0.2 - - if context.iteration_count < 5: - base_confidence += 0.1 - - return min(1.0, base_confidence) - - # Content creation methods (simplified implementations) - def _create_code_pattern_content( - self, content: str, patterns: List[str], context: ExtractionContext - ) -> str: - """Create formatted content for code pattern knowledge""" - return f"**Code Pattern: {', '.join(patterns)}**\n\n{content[:2000]}..." - - def _create_debugging_content(self, session: Dict) -> str: - """Create formatted content for debugging knowledge""" - return f"**Error:** {session.get('error_type', 'Unknown')}\n\n**Resolution:** {session.get('resolution', 'No resolution provided')}" - - def _create_process_content(self, context: ExtractionContext) -> str: - """Create formatted content for process knowledge""" - return f"**Task Type:** {context.task_data.get('task_type', 'Unknown')}\n**Iterations:** {context.iteration_count}\n**Duration:** {context.total_duration_minutes:.1f} minutes\n**Success:** {'Yes' if context.success else 'No'}" - - def _create_error_pattern_content( - self, pattern: Dict, context: ExtractionContext - ) -> str: - """Create formatted content for error pattern knowledge""" - return f"**Error Type:** {pattern['error_type']}\n**Frequency:** {pattern['frequency']}\n**Prevention:** {', '.join(pattern.get('prevention', []))}" - - def _create_optimization_content(self, context: ExtractionContext) -> str: - """Create formatted content for optimization knowledge""" - return f"**Optimization achieved in {context.total_duration_minutes:.1f} minutes with {context.iteration_count} iterations**" - - # Placeholder methods for more complex analysis functions - def _identify_debugging_sessions( - self, conversation_history: List[Dict] - ) -> List[Dict]: - """Identify debugging sessions in conversation history""" - return [] # Simplified implementation - - def _identify_error_patterns(self, conversation_history: List[Dict]) -> List[Dict]: - """Identify error patterns in conversation history""" - return [] # Simplified implementation - - def _identify_optimization_techniques( - self, context: ExtractionContext - ) -> List[str]: - """Identify optimization techniques used""" - return [] # Simplified implementation - - def _classify_problem_type(self, problem: str) -> str: - """Classify the type of problem""" - return "general" # Simplified implementation - - def _classify_solution_type(self, solution: str) -> str: - """Classify the type of solution""" - return "general" # Simplified implementation +""" +Task Knowledge Extractor for FuzeAgent + +This module extracts valuable knowledge from completed tasks and feeds it +into the hierarchical knowledge management system. It analyzes task outcomes, +code changes, conversation patterns, and performance metrics to create +reusable organizational knowledge. +""" + +import asyncio +import json +import logging +import re +from dataclasses import dataclass +from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple + +import asyncpg +from sentence_transformers import SentenceTransformer + +from .knowledge_propagation_engine import KnowledgePropagationEngine, PropagationTrigger +from .organization_rag_manager import ( + ContentType, + KnowledgeCategory, + OrganizationRAGManager, + SourceType, +) +from .team_knowledge_manager import TeamKnowledgeManager + +logger = logging.getLogger(__name__) + + +@dataclass +class TaskKnowledgeExtract: + """Represents extracted knowledge from a task""" + + title: str + content: str + content_type: ContentType + category: KnowledgeCategory + confidence_score: float + tags: List[str] + metadata: Dict[str, Any] + success_indicators: List[str] + failure_patterns: List[str] + + +@dataclass +class ExtractionContext: + """Context for knowledge extraction""" + + task_id: str + agent_id: str + team_id: str + organization_id: str + task_data: Dict[str, Any] + execution_result: Dict[str, Any] + conversation_history: List[Dict[str, Any]] + code_changes: List[Dict[str, Any]] + performance_metrics: Dict[str, Any] + iteration_count: int + total_duration_minutes: float + success: bool + + +class TaskKnowledgeExtractor: + """ + Extracts knowledge from completed tasks and integrates it + into the hierarchical knowledge management system. + """ + + def __init__( + self, + database_url: str, + org_rag_manager: OrganizationRAGManager, + team_knowledge_manager: TeamKnowledgeManager, + propagation_engine: KnowledgePropagationEngine, + ): + self.database_url = database_url + self.org_rag_manager = org_rag_manager + self.team_knowledge_manager = team_knowledge_manager + self.propagation_engine = propagation_engine + self.pool: Optional[asyncpg.Pool] = None + + # Initialize text analysis model + self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") + + # Extraction patterns and rules + self.code_patterns = self._initialize_code_patterns() + self.success_patterns = self._initialize_success_patterns() + self.failure_patterns = self._initialize_failure_patterns() + + # Configuration + self.min_extraction_confidence = 0.4 + self.min_task_duration_minutes = 5 # Don't extract from very short tasks + self.max_content_length = 5000 + + # Statistics + self.extractions_performed = 0 + self.knowledge_items_created = 0 + self.propagations_triggered = 0 + + async def initialize(self): + """Initialize the knowledge extractor""" + logger.info("Initializing TaskKnowledgeExtractor") + + try: + self.pool = await asyncpg.create_pool( + self.database_url, min_size=1, max_size=5, command_timeout=60 + ) + + logger.info("TaskKnowledgeExtractor initialized successfully") + + except Exception as e: + logger.error(f"Failed to initialize TaskKnowledgeExtractor: {e}") + raise + + async def close(self): + """Close database connections""" + if self.pool: + await self.pool.close() + logger.info("TaskKnowledgeExtractor closed") + + async def extract_knowledge_from_task( + self, task_id: str, agent_id: str, execution_result: Dict[str, Any] + ) -> List[str]: + """Extract knowledge from a completed task and store it""" + + try: + # Build extraction context + context = await self._build_extraction_context( + task_id, agent_id, execution_result + ) + + if not context: + logger.warning(f"Could not build extraction context for task {task_id}") + return [] + + # Skip extraction for very short or trivial tasks + if context.total_duration_minutes < self.min_task_duration_minutes: + logger.debug( + f"Skipping extraction for short task {task_id} ({context.total_duration_minutes:.1f}m)" + ) + return [] + + # Extract knowledge items + knowledge_extracts = await self._extract_knowledge_items(context) + + if not knowledge_extracts: + logger.debug(f"No knowledge extracted from task {task_id}") + return [] + + # Store extracted knowledge + stored_knowledge_ids = [] + for extract in knowledge_extracts: + if extract.confidence_score >= self.min_extraction_confidence: + knowledge_id = await self._store_knowledge_extract(context, extract) + if knowledge_id: + stored_knowledge_ids.append(knowledge_id) + + # Trigger knowledge propagation if we have valuable knowledge + if stored_knowledge_ids: + propagation_ids = ( + await self.propagation_engine.trigger_agent_to_team_propagation( + agent_id=context.agent_id, + task_id=context.task_id, + task_outcome={ + "success": context.success, + "task_type": context.task_data.get("task_type", "unknown"), + "complexity": self._assess_task_complexity(context), + "duration_minutes": context.total_duration_minutes, + "knowledge_extracted": len(stored_knowledge_ids), + "iteration_count": context.iteration_count, + }, + ) + ) + + self.propagations_triggered += len(propagation_ids) + logger.info( + f"Triggered {len(propagation_ids)} knowledge propagations for task {task_id}" + ) + + self.extractions_performed += 1 + self.knowledge_items_created += len(stored_knowledge_ids) + + logger.info( + f"Extracted {len(stored_knowledge_ids)} knowledge items from task {task_id}" + ) + return stored_knowledge_ids + + except Exception as e: + logger.error(f"Error extracting knowledge from task {task_id}: {e}") + return [] + + async def _build_extraction_context( + self, task_id: str, agent_id: str, execution_result: Dict[str, Any] + ) -> Optional[ExtractionContext]: + """Build context for knowledge extraction""" + + async with self.pool.acquire() as conn: + # Get basic task information + task_data = await conn.fetchrow( + """ + SELECT t.*, a.team_id, te.organization_id + FROM tasks t + JOIN agents a ON t.agent_id = a.id + JOIN teams te ON a.team_id = te.id + WHERE t.id = $1 + """, + task_id, + ) + + if not task_data: + return None + + # Get conversation history + conversation_history = await conn.fetch( + """ + SELECT message_type, content, metadata, created_at + FROM claude_conversations + WHERE task_id = $1 + ORDER BY created_at ASC + """, + task_id, + ) + + # Get code generations + code_changes = await conn.fetch( + """ + SELECT file_path, file_type, language, content, test_results, quality_metrics + FROM code_generations + WHERE task_id = $1 + ORDER BY generated_at ASC + """, + task_id, + ) + + # Get performance metrics + performance_metrics = await conn.fetch( + """ + SELECT metric_type, metric_value, metric_unit, context + FROM agent_performance_metrics + WHERE task_id = $1 + """, + task_id, + ) + + # Calculate duration + started_at = task_data["started_at"] + completed_at = execution_result.get("completed_at") + if completed_at: + if isinstance(completed_at, str): + completed_at = datetime.fromisoformat( + completed_at.replace("Z", "+00:00") + ) + duration = (completed_at - started_at).total_seconds() / 60.0 + else: + duration = 0.0 + + return ExtractionContext( + task_id=str(task_data["id"]), + agent_id=str(task_data["agent_id"]), + team_id=str(task_data["team_id"]), + organization_id=str(task_data["organization_id"]), + task_data=dict(task_data), + execution_result=execution_result, + conversation_history=[dict(conv) for conv in conversation_history], + code_changes=[dict(code) for code in code_changes], + performance_metrics={ + pm["metric_type"]: pm for pm in performance_metrics + }, + iteration_count=execution_result.get("iterations", 0), + total_duration_minutes=duration, + success=execution_result.get("status") == "completed", + ) + + async def _extract_knowledge_items( + self, context: ExtractionContext + ) -> List[TaskKnowledgeExtract]: + """Extract specific knowledge items from the task context""" + + knowledge_extracts = [] + + # Extract different types of knowledge + knowledge_extracts.extend(await self._extract_code_patterns(context)) + knowledge_extracts.extend(await self._extract_problem_solutions(context)) + knowledge_extracts.extend(await self._extract_debugging_insights(context)) + knowledge_extracts.extend(await self._extract_process_knowledge(context)) + knowledge_extracts.extend(await self._extract_error_patterns(context)) + knowledge_extracts.extend(await self._extract_optimization_insights(context)) + + return knowledge_extracts + + async def _extract_code_patterns( + self, context: ExtractionContext + ) -> List[TaskKnowledgeExtract]: + """Extract reusable code patterns and best practices""" + + extracts = [] + + for code_change in context.code_changes: + if code_change["file_type"] == "implementation": + content = code_change["content"] + language = code_change.get("language", "unknown") + + # Look for reusable patterns + patterns_found = [] + for pattern_name, pattern_info in self.code_patterns.items(): + if any( + keyword in content.lower() + for keyword in pattern_info["keywords"] + ): + patterns_found.append(pattern_name) + + if patterns_found and len(content) > 100: # Substantial code + # Create knowledge extract + title = f"Code Pattern: {', '.join(patterns_found)} ({language})" + extract_content = self._create_code_pattern_content( + content, patterns_found, context + ) + + confidence = self._calculate_code_pattern_confidence( + content, patterns_found, context + ) + + if confidence >= self.min_extraction_confidence: + extract = TaskKnowledgeExtract( + title=title, + content=extract_content, + content_type=ContentType.CODE, + category=KnowledgeCategory.DEVELOPMENT, + confidence_score=confidence, + tags=["code_pattern", language, *patterns_found], + metadata={ + "language": language, + "file_path": code_change["file_path"], + "patterns": patterns_found, + "task_success": context.success, + "lines_of_code": len(content.split("\n")), + }, + success_indicators=self._extract_success_indicators( + context + ), + failure_patterns=[], + ) + + extracts.append(extract) + + return extracts + + async def _extract_problem_solutions( + self, context: ExtractionContext + ) -> List[TaskKnowledgeExtract]: + """Extract problem-solution pairs from the task""" + + extracts = [] + + # Analyze conversation for problem descriptions and solutions + problem_solution_pairs = self._identify_problem_solution_pairs( + context.conversation_history + ) + + for problem, solution in problem_solution_pairs: + if len(problem) > 50 and len(solution) > 50: # Substantial content + title = f"Solution: {problem[:50]}..." + content = f"**Problem:**\n{problem}\n\n**Solution:**\n{solution}" + + # Determine category based on content + category = self._categorize_problem_solution(problem, solution) + + confidence = self._calculate_solution_confidence( + problem, solution, context + ) + + if confidence >= self.min_extraction_confidence: + extract = TaskKnowledgeExtract( + title=title, + content=content[: self.max_content_length], + content_type=ContentType.PROCEDURE, + category=category, + confidence_score=confidence, + tags=["problem_solution", "troubleshooting"], + metadata={ + "problem_type": self._classify_problem_type(problem), + "solution_type": self._classify_solution_type(solution), + "task_success": context.success, + }, + success_indicators=self._extract_success_indicators(context), + failure_patterns=[], + ) + + extracts.append(extract) + + return extracts + + async def _extract_debugging_insights( + self, context: ExtractionContext + ) -> List[TaskKnowledgeExtract]: + """Extract debugging approaches and insights""" + + extracts = [] + + # Look for error messages and resolution patterns + debugging_sessions = self._identify_debugging_sessions( + context.conversation_history + ) + + for session in debugging_sessions: + if session["resolution"] and context.success: + title = f"Debugging: {session['error_type']}" + content = self._create_debugging_content(session) + + confidence = self._calculate_debugging_confidence(session, context) + + if confidence >= self.min_extraction_confidence: + extract = TaskKnowledgeExtract( + title=title, + content=content, + content_type=ContentType.PROCEDURE, + category=KnowledgeCategory.TROUBLESHOOTING, + confidence_score=confidence, + tags=["debugging", session["error_type"], "troubleshooting"], + metadata={ + "error_type": session["error_type"], + "resolution_method": session["resolution_method"], + "tools_used": session.get("tools_used", []), + }, + success_indicators=self._extract_success_indicators(context), + failure_patterns=session.get("failure_patterns", []), + ) + + extracts.append(extract) + + return extracts + + async def _extract_process_knowledge( + self, context: ExtractionContext + ) -> List[TaskKnowledgeExtract]: + """Extract process and workflow knowledge""" + + extracts = [] + + if context.iteration_count > 1: # Multi-iteration tasks have process insights + title = f"Process: {context.task_data.get('task_type', 'Task')} Workflow" + + process_content = self._create_process_content(context) + confidence = self._calculate_process_confidence(context) + + if confidence >= self.min_extraction_confidence: + extract = TaskKnowledgeExtract( + title=title, + content=process_content, + content_type=ContentType.PROCEDURE, + category=KnowledgeCategory.PROCESS, + confidence_score=confidence, + tags=[ + "process", + "workflow", + context.task_data.get("task_type", "general"), + ], + metadata={ + "iterations_used": context.iteration_count, + "duration_minutes": context.total_duration_minutes, + "success_rate": 1.0 if context.success else 0.0, + "complexity": self._assess_task_complexity(context), + }, + success_indicators=self._extract_success_indicators(context), + failure_patterns=[], + ) + + extracts.append(extract) + + return extracts + + async def _extract_error_patterns( + self, context: ExtractionContext + ) -> List[TaskKnowledgeExtract]: + """Extract error patterns and avoidance strategies""" + + extracts = [] + + # Look for error patterns in failed tasks or recovered errors + error_patterns = self._identify_error_patterns(context.conversation_history) + + for pattern in error_patterns: + if pattern["frequency"] >= 2 or pattern["severity"] == "high": + title = f"Error Pattern: {pattern['error_type']}" + content = self._create_error_pattern_content(pattern, context) + + confidence = self._calculate_error_pattern_confidence(pattern, context) + + if confidence >= self.min_extraction_confidence: + extract = TaskKnowledgeExtract( + title=title, + content=content, + content_type=ContentType.DOCUMENTATION, + category=KnowledgeCategory.TROUBLESHOOTING, + confidence_score=confidence, + tags=["error_pattern", pattern["error_type"], "prevention"], + metadata={ + "error_type": pattern["error_type"], + "frequency": pattern["frequency"], + "severity": pattern["severity"], + "prevention_strategies": pattern.get("prevention", []), + }, + success_indicators=[], + failure_patterns=pattern.get("indicators", []), + ) + + extracts.append(extract) + + return extracts + + async def _extract_optimization_insights( + self, context: ExtractionContext + ) -> List[TaskKnowledgeExtract]: + """Extract performance optimization insights""" + + extracts = [] + + # Look for performance improvements in metrics + if "execution_time_minutes" in context.performance_metrics: + perf_data = context.performance_metrics["execution_time_minutes"] + if ( + perf_data["metric_value"] < 30 and context.success + ): # Efficient completion + title = "Performance Optimization: Efficient Task Execution" + content = self._create_optimization_content(context) + + confidence = self._calculate_optimization_confidence(context) + + if confidence >= self.min_extraction_confidence: + extract = TaskKnowledgeExtract( + title=title, + content=content, + content_type=ContentType.BEST_PRACTICE, + category=KnowledgeCategory.DEVELOPMENT, + confidence_score=confidence, + tags=["optimization", "performance", "efficiency"], + metadata={ + "execution_time": perf_data["metric_value"], + "iteration_efficiency": context.iteration_count + / context.total_duration_minutes, + "optimization_techniques": self._identify_optimization_techniques( + context + ), + }, + success_indicators=self._extract_success_indicators(context), + failure_patterns=[], + ) + + extracts.append(extract) + + return extracts + + async def _store_knowledge_extract( + self, context: ExtractionContext, extract: TaskKnowledgeExtract + ) -> Optional[str]: + """Store a knowledge extract in the appropriate knowledge base""" + + try: + # Store in organization knowledge base + knowledge_id = await self.org_rag_manager.add_knowledge( + organization_id=context.organization_id, + title=extract.title, + content=extract.content, + content_type=extract.content_type, + knowledge_category=extract.category, + source_type=SourceType.TASK_OUTCOME, + source_agent_id=context.agent_id, + source_team_id=context.team_id, + source_task_id=context.task_id, + relevance_score=extract.confidence_score, + quality_score=extract.confidence_score, + metadata={ + **extract.metadata, + "extraction_timestamp": datetime.now().isoformat(), + "extractor_version": "1.0", + "success_indicators": extract.success_indicators, + "failure_patterns": extract.failure_patterns, + }, + tags=extract.tags, + ) + + return knowledge_id + + except Exception as e: + logger.error(f"Error storing knowledge extract: {e}") + return None + + # Helper methods for pattern matching and analysis + def _initialize_code_patterns(self) -> Dict[str, Dict[str, Any]]: + """Initialize code pattern definitions""" + return { + "api_integration": { + "keywords": ["fetch", "axios", "request", "api", "endpoint", "rest"], + "confidence_boost": 0.2, + }, + "database_operations": { + "keywords": [ + "select", + "insert", + "update", + "delete", + "query", + "database", + "sql", + ], + "confidence_boost": 0.2, + }, + "authentication": { + "keywords": ["auth", "login", "token", "jwt", "session", "passport"], + "confidence_boost": 0.15, + }, + "error_handling": { + "keywords": ["try", "catch", "error", "exception", "throw"], + "confidence_boost": 0.1, + }, + "testing": { + "keywords": ["test", "spec", "describe", "it", "expect", "mock"], + "confidence_boost": 0.15, + }, + "optimization": { + "keywords": ["performance", "optimize", "cache", "memory", "speed"], + "confidence_boost": 0.2, + }, + } + + def _initialize_success_patterns(self) -> List[str]: + """Initialize success indicator patterns""" + return [ + r"test.*pass", + r"build.*success", + r"deploy.*complete", + r"fix.*issue", + r"resolve.*problem", + r"implement.*feature", + r"complete.*task", + ] + + def _initialize_failure_patterns(self) -> List[str]: + """Initialize failure indicator patterns""" + return [ + r"error.*occur", + r"fail.*to", + r"timeout.*exceed", + r"connection.*refuse", + r"not.*found", + r"access.*deni", + r"invalid.*request", + ] + + def _assess_task_complexity(self, context: ExtractionContext) -> str: + """Assess task complexity based on various factors""" + + complexity_score = 0 + + # Factor 1: Iteration count + if context.iteration_count > 10: + complexity_score += 3 + elif context.iteration_count > 5: + complexity_score += 2 + elif context.iteration_count > 2: + complexity_score += 1 + + # Factor 2: Duration + if context.total_duration_minutes > 180: # 3 hours + complexity_score += 3 + elif context.total_duration_minutes > 60: # 1 hour + complexity_score += 2 + elif context.total_duration_minutes > 30: + complexity_score += 1 + + # Factor 3: Code changes + if len(context.code_changes) > 10: + complexity_score += 2 + elif len(context.code_changes) > 5: + complexity_score += 1 + + # Factor 4: Conversation length + if len(context.conversation_history) > 50: + complexity_score += 2 + elif len(context.conversation_history) > 20: + complexity_score += 1 + + if complexity_score >= 6: + return "very_high" + elif complexity_score >= 4: + return "high" + elif complexity_score >= 2: + return "medium" + else: + return "low" + + def _extract_success_indicators(self, context: ExtractionContext) -> List[str]: + """Extract success indicators from the task execution""" + + indicators = [] + + # Look for success patterns in conversation + for conv in context.conversation_history: + content = conv.get("content", "").lower() + for pattern in self.success_patterns: + if re.search(pattern, content): + indicators.append(pattern) + + # Add task-specific indicators + if context.success: + indicators.append("task_completed_successfully") + + if context.execution_result.get("pull_request_url"): + indicators.append("pull_request_created") + + return list(set(indicators)) # Remove duplicates + + # Additional helper methods would be implemented here... + # (The file is getting quite long, so I'll implement key methods and indicate where others would go) + + def _identify_problem_solution_pairs( + self, conversation_history: List[Dict] + ) -> List[Tuple[str, str]]: + """Identify problem-solution pairs in conversation history""" + pairs = [] + # Implementation would analyze conversation flow to identify problems and their solutions + # This is a simplified placeholder + return pairs + + def _categorize_problem_solution( + self, problem: str, solution: str + ) -> KnowledgeCategory: + """Categorize a problem-solution pair""" + # Simple categorization based on keywords + combined_text = (problem + " " + solution).lower() + + if any(word in combined_text for word in ["test", "testing", "spec"]): + return KnowledgeCategory.TESTING + elif any(word in combined_text for word in ["deploy", "build", "ci", "cd"]): + return KnowledgeCategory.INFRASTRUCTURE + elif any(word in combined_text for word in ["security", "auth", "permission"]): + return KnowledgeCategory.SECURITY + elif any(word in combined_text for word in ["design", "ui", "ux", "interface"]): + return KnowledgeCategory.DESIGN + else: + return KnowledgeCategory.DEVELOPMENT + + def _calculate_code_pattern_confidence( + self, content: str, patterns: List[str], context: ExtractionContext + ) -> float: + """Calculate confidence score for code pattern extraction""" + base_confidence = 0.5 + + # Boost for successful task + if context.success: + base_confidence += 0.2 + + # Boost for multiple patterns + if len(patterns) > 1: + base_confidence += 0.1 + + # Boost for substantial code + if len(content) > 500: + base_confidence += 0.1 + + return min(1.0, base_confidence) + + def _calculate_solution_confidence( + self, problem: str, solution: str, context: ExtractionContext + ) -> float: + """Calculate confidence score for solution extraction""" + base_confidence = 0.4 + + if context.success: + base_confidence += 0.3 + + if len(solution) > 200: # Detailed solution + base_confidence += 0.1 + + return min(1.0, base_confidence) + + def _calculate_debugging_confidence( + self, session: Dict, context: ExtractionContext + ) -> float: + """Calculate confidence for debugging insights""" + base_confidence = 0.6 if context.success else 0.3 + + if session.get("resolution_method"): + base_confidence += 0.2 + + return min(1.0, base_confidence) + + def _calculate_process_confidence(self, context: ExtractionContext) -> float: + """Calculate confidence for process knowledge""" + if not context.success: + return 0.2 + + # Base confidence increases with iteration count (more process learning) + base_confidence = min(0.8, 0.3 + (context.iteration_count * 0.05)) + + return base_confidence + + def _calculate_error_pattern_confidence( + self, pattern: Dict, context: ExtractionContext + ) -> float: + """Calculate confidence for error pattern extraction""" + base_confidence = 0.4 + + if pattern["frequency"] > 2: + base_confidence += 0.2 + + if pattern["severity"] == "high": + base_confidence += 0.2 + + return min(1.0, base_confidence) + + def _calculate_optimization_confidence(self, context: ExtractionContext) -> float: + """Calculate confidence for optimization insights""" + if not context.success: + return 0.1 + + base_confidence = 0.5 + + # Boost for efficient execution + if context.total_duration_minutes < 30: + base_confidence += 0.2 + + if context.iteration_count < 5: + base_confidence += 0.1 + + return min(1.0, base_confidence) + + # Content creation methods (simplified implementations) + def _create_code_pattern_content( + self, content: str, patterns: List[str], context: ExtractionContext + ) -> str: + """Create formatted content for code pattern knowledge""" + return f"**Code Pattern: {', '.join(patterns)}**\n\n{content[:2000]}..." + + def _create_debugging_content(self, session: Dict) -> str: + """Create formatted content for debugging knowledge""" + return f"**Error:** {session.get('error_type', 'Unknown')}\n\n**Resolution:** {session.get('resolution', 'No resolution provided')}" + + def _create_process_content(self, context: ExtractionContext) -> str: + """Create formatted content for process knowledge""" + return f"**Task Type:** {context.task_data.get('task_type', 'Unknown')}\n**Iterations:** {context.iteration_count}\n**Duration:** {context.total_duration_minutes:.1f} minutes\n**Success:** {'Yes' if context.success else 'No'}" + + def _create_error_pattern_content( + self, pattern: Dict, context: ExtractionContext + ) -> str: + """Create formatted content for error pattern knowledge""" + return f"**Error Type:** {pattern['error_type']}\n**Frequency:** {pattern['frequency']}\n**Prevention:** {', '.join(pattern.get('prevention', []))}" + + def _create_optimization_content(self, context: ExtractionContext) -> str: + """Create formatted content for optimization knowledge""" + return f"**Optimization achieved in {context.total_duration_minutes:.1f} minutes with {context.iteration_count} iterations**" + + # Placeholder methods for more complex analysis functions + def _identify_debugging_sessions( + self, conversation_history: List[Dict] + ) -> List[Dict]: + """Identify debugging sessions in conversation history""" + return [] # Simplified implementation + + def _identify_error_patterns(self, conversation_history: List[Dict]) -> List[Dict]: + """Identify error patterns in conversation history""" + return [] # Simplified implementation + + def _identify_optimization_techniques( + self, context: ExtractionContext + ) -> List[str]: + """Identify optimization techniques used""" + return [] # Simplified implementation + + def _classify_problem_type(self, problem: str) -> str: + """Classify the type of problem""" + return "general" # Simplified implementation + + def _classify_solution_type(self, solution: str) -> str: + """Classify the type of solution""" + return "general" # Simplified implementation diff --git a/services/orchestrator/task_queue.py b/services/orchestrator/task_queue.py index 87ece7d..9178a92 100644 --- a/services/orchestrator/task_queue.py +++ b/services/orchestrator/task_queue.py @@ -1,124 +1,124 @@ -import asyncio -import json -import os -from typing import Any, Dict, List, Optional - -import aio_pika - -from .database import DatabaseManager - - -class TaskQueue: - def __init__(self): - self.rabbitmq_url = os.getenv( - "RABBITMQ_URL", "amqp://admin:password@rabbitmq:5672/" - ) - self.connection = None - self.channel = None - self.task_execution_engine = None # Will be set by orchestrator - - async def connect(self): - """Connect to RabbitMQ""" - if not self.connection: - self.connection = await aio_pika.connect_robust(self.rabbitmq_url) - self.channel = await self.connection.channel() - - async def assign_task(self, agent_id: str, task: dict) -> str: - """Assign a task to an agent""" - await self.connect() - - # Insert task into database - task_id = await DatabaseManager.insert_task( - title=task.get("title", "Untitled Task"), - description=task.get("description", ""), - assigned_to=agent_id, - created_by=task.get("created_by"), - ) - - # Add task_id to task data - task["id"] = task_id - task["assigned_to"] = agent_id - - # Send task to agent's queue - queue_name = f"agent_{agent_id.replace('-', '_')}" - queue = await self.channel.declare_queue(queue_name, durable=True) - - await self.channel.default_exchange.publish( - aio_pika.Message( - json.dumps(task).encode(), - delivery_mode=aio_pika.DeliveryMode.PERSISTENT, - ), - routing_key=queue_name, - ) - - return task_id - - async def list_tasks(self) -> List[Dict]: - """List all tasks""" - return await DatabaseManager.get_tasks() - - async def get_task(self, task_id: str) -> Dict: - """Get specific task""" - tasks = await self.list_tasks() - for task in tasks: - if str(task["id"]) == task_id: - return task - return None - - async def update_task_status(self, task_id: str, status: str, result: dict = None): - """Update task status""" - await DatabaseManager.update_task_status(task_id, status, result) - - async def get_pending_tasks(self) -> List[Dict]: - """Get all pending tasks""" - tasks = await self.list_tasks() - return [task for task in tasks if task["status"] == "pending"] - - async def get_agent_tasks(self, agent_id: str) -> List[Dict]: - """Get tasks assigned to specific agent""" - tasks = await self.list_tasks() - return [task for task in tasks if str(task["assigned_to"]) == agent_id] - - async def start_autonomous_execution(self, task_id: str) -> Dict[str, Any]: - """Start autonomous execution of a task""" - if not self.task_execution_engine: - raise RuntimeError("TaskExecutionEngine not configured") - - return await self.task_execution_engine.start_task_execution(task_id) - - async def get_execution_status(self, task_id: str) -> Dict[str, Any]: - """Get execution status of a task""" - if not self.task_execution_engine: - raise RuntimeError("TaskExecutionEngine not configured") - - return await self.task_execution_engine.get_execution_status(task_id) - - async def get_task_iterations(self, task_id: str) -> List[Dict[str, Any]]: - """Get task iteration history""" - if not self.task_execution_engine: - raise RuntimeError("TaskExecutionEngine not configured") - - return await self.task_execution_engine.get_task_iterations(task_id) - - async def handle_human_response(self, task_id: str, response: str) -> bool: - """Handle human response to a task question""" - if not self.task_execution_engine: - raise RuntimeError("TaskExecutionEngine not configured") - - return await self.task_execution_engine.handle_human_response(task_id, response) - - async def cancel_task_execution(self, task_id: str) -> bool: - """Cancel autonomous execution of a task""" - if not self.task_execution_engine: - raise RuntimeError("TaskExecutionEngine not configured") - - return await self.task_execution_engine.cancel_task_execution(task_id) - - def set_task_execution_engine(self, engine): - """Set the task execution engine reference""" - self.task_execution_engine = engine - - async def close(self): - """Close RabbitMQ connection""" - if self.connection: - await self.connection.close() +import asyncio +import json +import os +from typing import Any, Dict, List, Optional + +import aio_pika + +from .database import DatabaseManager + + +class TaskQueue: + def __init__(self): + self.rabbitmq_url = os.getenv( + "RABBITMQ_URL", "amqp://admin:password@rabbitmq:5672/" + ) + self.connection = None + self.channel = None + self.task_execution_engine = None # Will be set by orchestrator + + async def connect(self): + """Connect to RabbitMQ""" + if not self.connection: + self.connection = await aio_pika.connect_robust(self.rabbitmq_url) + self.channel = await self.connection.channel() + + async def assign_task(self, agent_id: str, task: dict) -> str: + """Assign a task to an agent""" + await self.connect() + + # Insert task into database + task_id = await DatabaseManager.insert_task( + title=task.get("title", "Untitled Task"), + description=task.get("description", ""), + assigned_to=agent_id, + created_by=task.get("created_by"), + ) + + # Add task_id to task data + task["id"] = task_id + task["assigned_to"] = agent_id + + # Send task to agent's queue + queue_name = f"agent_{agent_id.replace('-', '_')}" + queue = await self.channel.declare_queue(queue_name, durable=True) + + await self.channel.default_exchange.publish( + aio_pika.Message( + json.dumps(task).encode(), + delivery_mode=aio_pika.DeliveryMode.PERSISTENT, + ), + routing_key=queue_name, + ) + + return task_id + + async def list_tasks(self) -> List[Dict]: + """List all tasks""" + return await DatabaseManager.get_tasks() + + async def get_task(self, task_id: str) -> Dict: + """Get specific task""" + tasks = await self.list_tasks() + for task in tasks: + if str(task["id"]) == task_id: + return task + return None + + async def update_task_status(self, task_id: str, status: str, result: dict = None): + """Update task status""" + await DatabaseManager.update_task_status(task_id, status, result) + + async def get_pending_tasks(self) -> List[Dict]: + """Get all pending tasks""" + tasks = await self.list_tasks() + return [task for task in tasks if task["status"] == "pending"] + + async def get_agent_tasks(self, agent_id: str) -> List[Dict]: + """Get tasks assigned to specific agent""" + tasks = await self.list_tasks() + return [task for task in tasks if str(task["assigned_to"]) == agent_id] + + async def start_autonomous_execution(self, task_id: str) -> Dict[str, Any]: + """Start autonomous execution of a task""" + if not self.task_execution_engine: + raise RuntimeError("TaskExecutionEngine not configured") + + return await self.task_execution_engine.start_task_execution(task_id) + + async def get_execution_status(self, task_id: str) -> Dict[str, Any]: + """Get execution status of a task""" + if not self.task_execution_engine: + raise RuntimeError("TaskExecutionEngine not configured") + + return await self.task_execution_engine.get_execution_status(task_id) + + async def get_task_iterations(self, task_id: str) -> List[Dict[str, Any]]: + """Get task iteration history""" + if not self.task_execution_engine: + raise RuntimeError("TaskExecutionEngine not configured") + + return await self.task_execution_engine.get_task_iterations(task_id) + + async def handle_human_response(self, task_id: str, response: str) -> bool: + """Handle human response to a task question""" + if not self.task_execution_engine: + raise RuntimeError("TaskExecutionEngine not configured") + + return await self.task_execution_engine.handle_human_response(task_id, response) + + async def cancel_task_execution(self, task_id: str) -> bool: + """Cancel autonomous execution of a task""" + if not self.task_execution_engine: + raise RuntimeError("TaskExecutionEngine not configured") + + return await self.task_execution_engine.cancel_task_execution(task_id) + + def set_task_execution_engine(self, engine): + """Set the task execution engine reference""" + self.task_execution_engine = engine + + async def close(self): + """Close RabbitMQ connection""" + if self.connection: + await self.connection.close() diff --git a/services/orchestrator/team_knowledge_manager.py b/services/orchestrator/team_knowledge_manager.py index f7db4ab..a20741b 100644 --- a/services/orchestrator/team_knowledge_manager.py +++ b/services/orchestrator/team_knowledge_manager.py @@ -1,869 +1,869 @@ -""" -Team Knowledge Manager for FuzeAgent - -This module manages team-level knowledge aggregation, filtering organization knowledge -for team relevance, and facilitating knowledge sharing between agents within teams. -""" - -import asyncio -import json -import logging -import uuid -from dataclasses import dataclass -from datetime import datetime, timedelta -from typing import Any, Dict, List, Optional, Tuple - -import asyncpg -from sentence_transformers import SentenceTransformer - -from .organization_rag_manager import ( - ContentType, - KnowledgeCategory, - KnowledgeSearchResult, - OrganizationRAGManager, - SourceType, - VisibilityLevel, -) - -logger = logging.getLogger(__name__) - - -@dataclass -class TeamKnowledge: - """Represents team-level knowledge""" - - id: str - team_id: str - organization_id: str - title: str - content: str - content_type: ContentType - knowledge_category: KnowledgeCategory - embedding: Optional[List[float]] - source_type: SourceType - contributing_agents: List[str] - source_knowledge_ids: List[str] - aggregation_method: str - team_relevance_score: float - agent_adoption_rate: float - effectiveness_score: float - visibility_level: VisibilityLevel - metadata: Dict[str, Any] - tags: List[str] - created_at: datetime - updated_at: datetime - last_accessed: Optional[datetime] - - -@dataclass -class TeamKnowledgeSearchResult: - """Result of team knowledge search""" - - team_knowledge: TeamKnowledge - similarity_score: float - relevance_score: float - team_fit_score: float - combined_score: float - - -class TeamKnowledgeManager: - """ - Manages team-specific knowledge base with intelligent aggregation - from organization knowledge and agent contributions. - """ - - def __init__( - self, database_url: str, organization_rag_manager: OrganizationRAGManager - ): - self.database_url = database_url - self.org_rag_manager = organization_rag_manager - self.pool: Optional[asyncpg.Pool] = None - - # Initialize embedding model - self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") - self.embedding_dim = 384 - - # Configuration - self.min_team_relevance = 0.4 - self.adoption_threshold = 0.6 # 60% of team agents should find it useful - self.effectiveness_decay_days = 30 - - # Statistics - self.team_queries_processed = 0 - self.team_knowledge_created = 0 - self.aggregations_performed = 0 - - async def initialize(self): - """Initialize the team knowledge manager""" - logger.info("Initializing TeamKnowledgeManager") - - try: - self.pool = await asyncpg.create_pool( - self.database_url, min_size=2, max_size=10, command_timeout=60 - ) - - logger.info("TeamKnowledgeManager initialized successfully") - - except Exception as e: - logger.error(f"Failed to initialize TeamKnowledgeManager: {e}") - raise - - async def close(self): - """Close database connections""" - if self.pool: - await self.pool.close() - logger.info("TeamKnowledgeManager closed") - - async def create_team_knowledge( - self, - team_id: str, - title: str, - content: str, - content_type: ContentType = ContentType.TEXT, - knowledge_category: KnowledgeCategory = KnowledgeCategory.DEVELOPMENT, - source_type: SourceType = SourceType.TEAM_AGGREGATION, - contributing_agents: Optional[List[str]] = None, - source_knowledge_ids: Optional[List[str]] = None, - aggregation_method: str = "synthesis", - team_relevance_score: float = 0.7, - metadata: Optional[Dict[str, Any]] = None, - tags: Optional[List[str]] = None, - ) -> str: - """Create team-specific knowledge""" - - team_knowledge_id = str(uuid.uuid4()) - embedding = self._generate_embedding(content) - - async with self.pool.acquire() as conn: - # Get organization_id for this team - org_id = await conn.fetchval( - """ - SELECT organization_id FROM teams WHERE id = $1 - """, - team_id, - ) - - if not org_id: - raise ValueError(f"Team {team_id} not found") - - await conn.execute( - """ - INSERT INTO team_knowledge_base ( - id, team_id, organization_id, title, content, content_type, - knowledge_category, embedding, source_type, contributing_agents, - source_knowledge_ids, aggregation_method, team_relevance_score, - metadata, tags - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15) - """, - team_knowledge_id, - team_id, - org_id, - title, - content, - content_type.value, - knowledge_category.value, - embedding, - source_type.value, - contributing_agents or [], - source_knowledge_ids or [], - aggregation_method, - team_relevance_score, - json.dumps(metadata or {}), - tags or [], - ) - - self.team_knowledge_created += 1 - - logger.info(f"Created team knowledge {team_knowledge_id} for team {team_id}") - return team_knowledge_id - - async def search_team_knowledge( - self, - team_id: str, - query: str, - categories: Optional[List[KnowledgeCategory]] = None, - content_types: Optional[List[ContentType]] = None, - include_org_knowledge: bool = True, - limit: int = 10, - min_similarity: float = 0.3, - ) -> List[TeamKnowledgeSearchResult]: - """Search team knowledge with optional organization knowledge inclusion""" - - self.team_queries_processed += 1 - query_embedding = self._generate_embedding(query) - results = [] - - async with self.pool.acquire() as conn: - # Search team-specific knowledge - team_results = await self._search_team_specific_knowledge( - conn, - team_id, - query_embedding, - categories, - content_types, - limit, - min_similarity, - ) - results.extend(team_results) - - # Search organization knowledge filtered for team relevance - if include_org_knowledge and len(results) < limit: - org_results = await self._search_org_knowledge_for_team( - conn, - team_id, - query, - categories, - content_types, - limit - len(results), - min_similarity, - ) - results.extend(org_results) - - # Sort by combined score - results.sort(key=lambda x: x.combined_score, reverse=True) - return results[:limit] - - async def aggregate_agent_knowledge_to_team( - self, - team_id: str, - agent_id: str, - agent_memory_ids: List[str], - aggregation_method: str = "synthesis", - ) -> Optional[str]: - """Aggregate multiple agent memories into team knowledge""" - - async with self.pool.acquire() as conn: - # Get agent memories - agent_memories = await conn.fetch( - """ - SELECT * FROM agent_memory - WHERE id = ANY($1) AND agent_id = $2 - ORDER BY confidence_score DESC, created_at DESC - """, - agent_memory_ids, - agent_id, - ) - - if not agent_memories: - return None - - # Analyze memories for commonalities - analysis_result = await self._analyze_memories_for_aggregation( - agent_memories - ) - - if analysis_result["aggregation_value"] < self.min_team_relevance: - logger.debug( - f"Agent memories don't meet team relevance threshold: {analysis_result['aggregation_value']}" - ) - return None - - # Create aggregated knowledge - team_knowledge_id = await self.create_team_knowledge( - team_id=team_id, - title=analysis_result["title"], - content=analysis_result["content"], - content_type=analysis_result["content_type"], - knowledge_category=analysis_result["category"], - source_type=SourceType.AGENT_CONTRIBUTION, - contributing_agents=[agent_id], - aggregation_method=aggregation_method, - team_relevance_score=analysis_result["aggregation_value"], - metadata=analysis_result["metadata"], - tags=analysis_result["tags"], - ) - - # Mark original memories as aggregated - await conn.execute( - """ - UPDATE agent_memory - SET propagated_to_team = TRUE, team_context_id = $2 - WHERE id = ANY($1) - """, - agent_memory_ids, - team_knowledge_id, - ) - - self.aggregations_performed += 1 - - logger.info( - f"Aggregated {len(agent_memory_ids)} agent memories into team knowledge {team_knowledge_id}" - ) - return team_knowledge_id - - async def get_team_knowledge_context( - self, - team_id: str, - task_context: Dict[str, Any], - agent_id: Optional[str] = None, - max_context_items: int = 5, - ) -> Dict[str, Any]: - """Get relevant team knowledge for task execution context""" - - # Build context query from task information - context_query = self._build_context_query(task_context) - - # Search for relevant knowledge - search_results = await self.search_team_knowledge( - team_id=team_id, - query=context_query, - limit=max_context_items, - min_similarity=0.4, - ) - - # Get team statistics - team_stats = await self.get_team_knowledge_stats(team_id) - - # Build context - context = { - "team_id": team_id, - "relevant_knowledge": [ - { - "id": result.team_knowledge.id, - "title": result.team_knowledge.title, - "content": ( - result.team_knowledge.content[:500] + "..." - if len(result.team_knowledge.content) > 500 - else result.team_knowledge.content - ), - "category": result.team_knowledge.knowledge_category.value, - "relevance_score": result.combined_score, - "usage_stats": { - "adoption_rate": result.team_knowledge.agent_adoption_rate, - "effectiveness": result.team_knowledge.effectiveness_score, - }, - } - for result in search_results - ], - "team_knowledge_stats": team_stats, - "context_query": context_query, - "generated_at": datetime.now().isoformat(), - } - - return context - - async def update_knowledge_effectiveness( - self, - team_knowledge_id: str, - agent_id: str, - task_success: bool, - feedback_score: Optional[float] = None, - usage_context: Optional[Dict[str, Any]] = None, - ): - """Update knowledge effectiveness based on agent usage""" - - async with self.pool.acquire() as conn: - # Get current knowledge - knowledge = await conn.fetchrow( - """ - SELECT * FROM team_knowledge_base WHERE id = $1 - """, - team_knowledge_id, - ) - - if not knowledge: - return - - # Calculate new effectiveness score - success_weight = 1.0 if task_success else -0.3 - feedback_weight = (feedback_score or 0.5) - 0.5 - - # Update effectiveness with exponential moving average - current_effectiveness = knowledge["effectiveness_score"] - new_effectiveness = ( - current_effectiveness * 0.8 + (success_weight + feedback_weight) * 0.2 - ) - new_effectiveness = max(0.0, min(1.0, new_effectiveness)) - - # Update agent adoption tracking - contributing_agents = knowledge["contributing_agents"] or [] - if agent_id not in contributing_agents: - contributing_agents.append(agent_id) - - # Calculate adoption rate (agents who used it / total team agents) - team_agent_count = await conn.fetchval( - """ - SELECT COUNT(*) FROM agents WHERE team_id = $1 - """, - knowledge["team_id"], - ) - - adoption_rate = len(contributing_agents) / max(1, team_agent_count) - - # Update knowledge - await conn.execute( - """ - UPDATE team_knowledge_base - SET effectiveness_score = $2, - agent_adoption_rate = $3, - contributing_agents = $4, - last_accessed = NOW(), - updated_at = NOW() - WHERE id = $1 - """, - team_knowledge_id, - new_effectiveness, - adoption_rate, - contributing_agents, - ) - - logger.debug( - f"Updated knowledge {team_knowledge_id} effectiveness: {new_effectiveness:.2f}, adoption: {adoption_rate:.2f}" - ) - - async def get_team_knowledge_stats(self, team_id: str) -> Dict[str, Any]: - """Get comprehensive team knowledge statistics""" - - async with self.pool.acquire() as conn: - # Basic statistics - basic_stats = await conn.fetchrow( - """ - SELECT - COUNT(*) as total_knowledge, - COUNT(DISTINCT knowledge_category) as categories, - COUNT(DISTINCT unnest(contributing_agents)) as contributing_agents, - AVG(team_relevance_score) as avg_relevance, - AVG(effectiveness_score) as avg_effectiveness, - AVG(agent_adoption_rate) as avg_adoption_rate - FROM team_knowledge_base - WHERE team_id = $1 - """, - team_id, - ) - - # Category breakdown - category_stats = await conn.fetch( - """ - SELECT - knowledge_category, - COUNT(*) as count, - AVG(effectiveness_score) as avg_effectiveness, - AVG(agent_adoption_rate) as avg_adoption - FROM team_knowledge_base - WHERE team_id = $1 - GROUP BY knowledge_category - ORDER BY count DESC - """, - team_id, - ) - - # Most effective knowledge - top_knowledge = await conn.fetch( - """ - SELECT - title, - knowledge_category, - effectiveness_score, - agent_adoption_rate - FROM team_knowledge_base - WHERE team_id = $1 - ORDER BY effectiveness_score DESC - LIMIT 5 - """, - team_id, - ) - - return { - "team_id": team_id, - "basic_stats": dict(basic_stats) if basic_stats else {}, - "category_breakdown": [dict(cat) for cat in category_stats], - "top_knowledge": [dict(know) for know in top_knowledge], - "generated_at": datetime.now().isoformat(), - } - - async def _search_team_specific_knowledge( - self, - conn, - team_id: str, - query_embedding: List[float], - categories: Optional[List[KnowledgeCategory]], - content_types: Optional[List[ContentType]], - limit: int, - min_similarity: float, - ) -> List[TeamKnowledgeSearchResult]: - """Search team-specific knowledge base""" - - # Build query conditions - where_conditions = ["team_id = $2"] - params = [query_embedding, team_id] - param_idx = 3 - - if categories: - where_conditions.append(f"knowledge_category = ANY(${param_idx})") - params.append([cat.value for cat in categories]) - param_idx += 1 - - if content_types: - where_conditions.append(f"content_type = ANY(${param_idx})") - params.append([ct.value for ct in content_types]) - param_idx += 1 - - where_conditions.append(f"(1 - (embedding <=> $1)) >= ${param_idx}") - params.append(min_similarity) - param_idx += 1 - - where_clause = "WHERE " + " AND ".join(where_conditions) - - results = await conn.fetch( - f""" - SELECT - *, - (1 - (embedding <=> $1)) as similarity_score - FROM team_knowledge_base - {where_clause} - ORDER BY similarity_score DESC, effectiveness_score DESC - LIMIT ${param_idx} - """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params - *params, - limit, - ) - - search_results = [] - for row in results: - team_knowledge = self._row_to_team_knowledge(row) - - # Calculate team fit score based on adoption and effectiveness - team_fit_score = ( - team_knowledge.agent_adoption_rate * 0.4 - + team_knowledge.effectiveness_score * 0.6 - ) - - combined_score = ( - float(row["similarity_score"]) * 0.4 - + team_knowledge.team_relevance_score * 0.3 - + team_fit_score * 0.3 - ) - - search_results.append( - TeamKnowledgeSearchResult( - team_knowledge=team_knowledge, - similarity_score=float(row["similarity_score"]), - relevance_score=team_knowledge.team_relevance_score, - team_fit_score=team_fit_score, - combined_score=combined_score, - ) - ) - - return search_results - - async def _search_org_knowledge_for_team( - self, - conn, - team_id: str, - query: str, - categories: Optional[List[KnowledgeCategory]], - content_types: Optional[List[ContentType]], - limit: int, - min_similarity: float, - ) -> List[TeamKnowledgeSearchResult]: - """Search organization knowledge filtered for team relevance""" - - # Get organization ID for the team - org_id = await conn.fetchval( - """ - SELECT organization_id FROM teams WHERE id = $1 - """, - team_id, - ) - - if not org_id: - return [] - - # Search organization knowledge - org_results = await self.org_rag_manager.search_knowledge( - organization_id=str(org_id), - query=query, - categories=categories, - content_types=content_types, - limit=limit * 2, # Get more to filter for team relevance - min_similarity=min_similarity, - requester_team_id=team_id, - ) - - # Convert to team knowledge search results with team relevance scoring - team_results = [] - for org_result in org_results: - # Calculate team relevance based on source and usage - team_relevance = await self._calculate_team_relevance( - conn, team_id, org_result.knowledge - ) - - if team_relevance >= self.min_team_relevance: - # Create pseudo team knowledge for consistent interface - pseudo_team_knowledge = TeamKnowledge( - id=org_result.knowledge.id, - team_id=team_id, - organization_id=org_result.knowledge.organization_id, - title=org_result.knowledge.title, - content=org_result.knowledge.content, - content_type=org_result.knowledge.content_type, - knowledge_category=org_result.knowledge.knowledge_category, - embedding=org_result.knowledge.embedding, - source_type=org_result.knowledge.source_type, - contributing_agents=[], - source_knowledge_ids=[org_result.knowledge.id], - aggregation_method="organization_filter", - team_relevance_score=team_relevance, - agent_adoption_rate=0.0, - effectiveness_score=org_result.knowledge.success_correlation, - visibility_level=org_result.knowledge.visibility_level, - metadata=org_result.knowledge.metadata, - tags=org_result.knowledge.tags, - created_at=org_result.knowledge.created_at, - updated_at=org_result.knowledge.updated_at, - last_accessed=org_result.knowledge.last_accessed, - ) - - combined_score = ( - org_result.similarity_score * 0.5 + team_relevance * 0.5 - ) - - team_results.append( - TeamKnowledgeSearchResult( - team_knowledge=pseudo_team_knowledge, - similarity_score=org_result.similarity_score, - relevance_score=org_result.relevance_score, - team_fit_score=team_relevance, - combined_score=combined_score, - ) - ) - - return team_results[:limit] - - async def _calculate_team_relevance( - self, conn, team_id: str, org_knowledge - ) -> float: - """Calculate how relevant organization knowledge is for a specific team""" - - relevance_factors = [] - - # Factor 1: Source team match - if org_knowledge.source_team_id == team_id: - relevance_factors.append(1.0) - elif org_knowledge.source_team_id: - # Check if source team is similar to current team - team_similarity = await self._calculate_team_similarity( - conn, team_id, org_knowledge.source_team_id - ) - relevance_factors.append(team_similarity) - else: - relevance_factors.append(0.3) # No team context - - # Factor 2: Category relevance to team's work - team_categories = await self._get_team_primary_categories(conn, team_id) - if org_knowledge.knowledge_category.value in team_categories: - relevance_factors.append(0.9) - else: - relevance_factors.append(0.4) - - # Factor 3: Usage by team agents - team_usage = ( - await conn.fetchval( - """ - SELECT COUNT(DISTINCT source_agent_id)::float / NULLIF( - (SELECT COUNT(*) FROM agents WHERE team_id = $1), 0 - ) - FROM organization_knowledge_base - WHERE id = $2 AND source_agent_id IN ( - SELECT id FROM agents WHERE team_id = $1 - ) - """, - team_id, - org_knowledge.id, - ) - or 0.0 - ) - relevance_factors.append(team_usage) - - # Factor 4: Base quality and relevance - relevance_factors.append(org_knowledge.quality_score) - relevance_factors.append(org_knowledge.relevance_score) - - # Calculate weighted average - weights = [0.3, 0.25, 0.25, 0.1, 0.1] - team_relevance = sum(f * w for f, w in zip(relevance_factors, weights)) - - return min(1.0, max(0.0, team_relevance)) - - async def _calculate_team_similarity( - self, conn, team_id1: str, team_id2: str - ) -> float: - """Calculate similarity between two teams based on their work patterns""" - - # Simple implementation based on team type and settings - team_info = await conn.fetch( - """ - SELECT id, team_type, settings FROM teams - WHERE id IN ($1, $2) - """, - team_id1, - team_id2, - ) - - if len(team_info) != 2: - return 0.0 - - team1, team2 = team_info - - # Type similarity - type_similarity = 1.0 if team1["team_type"] == team2["team_type"] else 0.5 - - # Settings similarity (simplified) - settings1 = team1["settings"] or {} - settings2 = team2["settings"] or {} - - common_keys = set(settings1.keys()) & set(settings2.keys()) - if common_keys: - settings_similarity = sum( - 1.0 if settings1.get(key) == settings2.get(key) else 0.0 - for key in common_keys - ) / len(common_keys) - else: - settings_similarity = 0.5 - - return type_similarity * 0.7 + settings_similarity * 0.3 - - async def _get_team_primary_categories(self, conn, team_id: str) -> List[str]: - """Get primary knowledge categories this team works with""" - - categories = await conn.fetch( - """ - SELECT knowledge_category, COUNT(*) as usage_count - FROM team_knowledge_base - WHERE team_id = $1 - GROUP BY knowledge_category - ORDER BY usage_count DESC - LIMIT 3 - """, - team_id, - ) - - return [cat["knowledge_category"] for cat in categories] - - def _generate_embedding(self, text: str) -> List[float]: - """Generate embedding for text using sentence transformers""" - try: - embedding = self.embedding_model.encode(text, convert_to_tensor=False) - return embedding.tolist() - except Exception as e: - logger.error(f"Error generating embedding: {e}") - return [0.0] * self.embedding_dim - - def _row_to_team_knowledge(self, row) -> TeamKnowledge: - """Convert database row to TeamKnowledge object""" - return TeamKnowledge( - id=str(row["id"]), - team_id=str(row["team_id"]), - organization_id=str(row["organization_id"]), - title=row["title"], - content=row["content"], - content_type=ContentType(row["content_type"]), - knowledge_category=KnowledgeCategory(row["knowledge_category"]), - embedding=row["embedding"] if row["embedding"] else None, - source_type=SourceType(row["source_type"]), - contributing_agents=row["contributing_agents"] or [], - source_knowledge_ids=row["source_knowledge_ids"] or [], - aggregation_method=row["aggregation_method"], - team_relevance_score=row["team_relevance_score"], - agent_adoption_rate=row["agent_adoption_rate"], - effectiveness_score=row["effectiveness_score"], - visibility_level=VisibilityLevel(row["visibility_level"]), - metadata=( - json.loads(row["metadata"]) - if isinstance(row["metadata"], str) - else row["metadata"] - ), - tags=row["tags"] or [], - created_at=row["created_at"], - updated_at=row["updated_at"], - last_accessed=row["last_accessed"], - ) - - def _build_context_query(self, task_context: Dict[str, Any]) -> str: - """Build a search query from task context""" - query_parts = [] - - if task_context.get("task_type"): - query_parts.append(task_context["task_type"]) - - if task_context.get("description"): - query_parts.append(task_context["description"]) - - if task_context.get("technologies"): - query_parts.extend(task_context["technologies"]) - - if task_context.get("domain"): - query_parts.append(task_context["domain"]) - - return " ".join(query_parts) - - async def _analyze_memories_for_aggregation( - self, agent_memories: List - ) -> Dict[str, Any]: - """Analyze agent memories to determine if they should be aggregated""" - - if not agent_memories: - return {"aggregation_value": 0.0} - - # Simple aggregation analysis - # In practice, this could use more sophisticated NLP - - # Calculate average confidence and success correlation - avg_confidence = sum(mem["confidence_score"] for mem in agent_memories) / len( - agent_memories - ) - avg_success = sum( - mem.get("success_correlation", 0.0) for mem in agent_memories - ) / len(agent_memories) - - # Find common themes - all_content = " ".join(mem["content"] for mem in agent_memories) - - # Determine primary category - categories = [mem.get("memory_type", "general") for mem in agent_memories] - primary_category = ( - max(set(categories), key=categories.count) if categories else "general" - ) - - # Create aggregated content (simplified) - title = f"Team Knowledge: {primary_category.replace('_', ' ').title()}" - content = ( - f"Aggregated knowledge from {len(agent_memories)} agent experiences:\n\n" - + all_content[:1000] - ) - - # Map memory type to knowledge category - category_mapping = { - "code_pattern": KnowledgeCategory.DEVELOPMENT, - "task_outcome": KnowledgeCategory.PROCESS, - "debugging": KnowledgeCategory.TROUBLESHOOTING, - "optimization": KnowledgeCategory.DEVELOPMENT, - "testing": KnowledgeCategory.TESTING, - } - - knowledge_category = category_mapping.get( - primary_category, KnowledgeCategory.DEVELOPMENT - ) - - # Determine content type - content_type = ( - ContentType.CODE if "code" in primary_category else ContentType.TEXT - ) - - # Calculate aggregation value - aggregation_value = min(1.0, (avg_confidence + avg_success) / 2.0) - - return { - "aggregation_value": aggregation_value, - "title": title, - "content": content, - "content_type": content_type, - "category": knowledge_category, - "metadata": { - "source_memory_count": len(agent_memories), - "avg_confidence": avg_confidence, - "avg_success_correlation": avg_success, - "primary_type": primary_category, - }, - "tags": [primary_category, "aggregated", "agent_contribution"], - } +""" +Team Knowledge Manager for FuzeAgent + +This module manages team-level knowledge aggregation, filtering organization knowledge +for team relevance, and facilitating knowledge sharing between agents within teams. +""" + +import asyncio +import json +import logging +import uuid +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Any, Dict, List, Optional, Tuple + +import asyncpg +from sentence_transformers import SentenceTransformer + +from .organization_rag_manager import ( + ContentType, + KnowledgeCategory, + KnowledgeSearchResult, + OrganizationRAGManager, + SourceType, + VisibilityLevel, +) + +logger = logging.getLogger(__name__) + + +@dataclass +class TeamKnowledge: + """Represents team-level knowledge""" + + id: str + team_id: str + organization_id: str + title: str + content: str + content_type: ContentType + knowledge_category: KnowledgeCategory + embedding: Optional[List[float]] + source_type: SourceType + contributing_agents: List[str] + source_knowledge_ids: List[str] + aggregation_method: str + team_relevance_score: float + agent_adoption_rate: float + effectiveness_score: float + visibility_level: VisibilityLevel + metadata: Dict[str, Any] + tags: List[str] + created_at: datetime + updated_at: datetime + last_accessed: Optional[datetime] + + +@dataclass +class TeamKnowledgeSearchResult: + """Result of team knowledge search""" + + team_knowledge: TeamKnowledge + similarity_score: float + relevance_score: float + team_fit_score: float + combined_score: float + + +class TeamKnowledgeManager: + """ + Manages team-specific knowledge base with intelligent aggregation + from organization knowledge and agent contributions. + """ + + def __init__( + self, database_url: str, organization_rag_manager: OrganizationRAGManager + ): + self.database_url = database_url + self.org_rag_manager = organization_rag_manager + self.pool: Optional[asyncpg.Pool] = None + + # Initialize embedding model + self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") + self.embedding_dim = 384 + + # Configuration + self.min_team_relevance = 0.4 + self.adoption_threshold = 0.6 # 60% of team agents should find it useful + self.effectiveness_decay_days = 30 + + # Statistics + self.team_queries_processed = 0 + self.team_knowledge_created = 0 + self.aggregations_performed = 0 + + async def initialize(self): + """Initialize the team knowledge manager""" + logger.info("Initializing TeamKnowledgeManager") + + try: + self.pool = await asyncpg.create_pool( + self.database_url, min_size=2, max_size=10, command_timeout=60 + ) + + logger.info("TeamKnowledgeManager initialized successfully") + + except Exception as e: + logger.error(f"Failed to initialize TeamKnowledgeManager: {e}") + raise + + async def close(self): + """Close database connections""" + if self.pool: + await self.pool.close() + logger.info("TeamKnowledgeManager closed") + + async def create_team_knowledge( + self, + team_id: str, + title: str, + content: str, + content_type: ContentType = ContentType.TEXT, + knowledge_category: KnowledgeCategory = KnowledgeCategory.DEVELOPMENT, + source_type: SourceType = SourceType.TEAM_AGGREGATION, + contributing_agents: Optional[List[str]] = None, + source_knowledge_ids: Optional[List[str]] = None, + aggregation_method: str = "synthesis", + team_relevance_score: float = 0.7, + metadata: Optional[Dict[str, Any]] = None, + tags: Optional[List[str]] = None, + ) -> str: + """Create team-specific knowledge""" + + team_knowledge_id = str(uuid.uuid4()) + embedding = self._generate_embedding(content) + + async with self.pool.acquire() as conn: + # Get organization_id for this team + org_id = await conn.fetchval( + """ + SELECT organization_id FROM teams WHERE id = $1 + """, + team_id, + ) + + if not org_id: + raise ValueError(f"Team {team_id} not found") + + await conn.execute( + """ + INSERT INTO team_knowledge_base ( + id, team_id, organization_id, title, content, content_type, + knowledge_category, embedding, source_type, contributing_agents, + source_knowledge_ids, aggregation_method, team_relevance_score, + metadata, tags + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15) + """, + team_knowledge_id, + team_id, + org_id, + title, + content, + content_type.value, + knowledge_category.value, + embedding, + source_type.value, + contributing_agents or [], + source_knowledge_ids or [], + aggregation_method, + team_relevance_score, + json.dumps(metadata or {}), + tags or [], + ) + + self.team_knowledge_created += 1 + + logger.info(f"Created team knowledge {team_knowledge_id} for team {team_id}") + return team_knowledge_id + + async def search_team_knowledge( + self, + team_id: str, + query: str, + categories: Optional[List[KnowledgeCategory]] = None, + content_types: Optional[List[ContentType]] = None, + include_org_knowledge: bool = True, + limit: int = 10, + min_similarity: float = 0.3, + ) -> List[TeamKnowledgeSearchResult]: + """Search team knowledge with optional organization knowledge inclusion""" + + self.team_queries_processed += 1 + query_embedding = self._generate_embedding(query) + results = [] + + async with self.pool.acquire() as conn: + # Search team-specific knowledge + team_results = await self._search_team_specific_knowledge( + conn, + team_id, + query_embedding, + categories, + content_types, + limit, + min_similarity, + ) + results.extend(team_results) + + # Search organization knowledge filtered for team relevance + if include_org_knowledge and len(results) < limit: + org_results = await self._search_org_knowledge_for_team( + conn, + team_id, + query, + categories, + content_types, + limit - len(results), + min_similarity, + ) + results.extend(org_results) + + # Sort by combined score + results.sort(key=lambda x: x.combined_score, reverse=True) + return results[:limit] + + async def aggregate_agent_knowledge_to_team( + self, + team_id: str, + agent_id: str, + agent_memory_ids: List[str], + aggregation_method: str = "synthesis", + ) -> Optional[str]: + """Aggregate multiple agent memories into team knowledge""" + + async with self.pool.acquire() as conn: + # Get agent memories + agent_memories = await conn.fetch( + """ + SELECT * FROM agent_memory + WHERE id = ANY($1) AND agent_id = $2 + ORDER BY confidence_score DESC, created_at DESC + """, + agent_memory_ids, + agent_id, + ) + + if not agent_memories: + return None + + # Analyze memories for commonalities + analysis_result = await self._analyze_memories_for_aggregation( + agent_memories + ) + + if analysis_result["aggregation_value"] < self.min_team_relevance: + logger.debug( + f"Agent memories don't meet team relevance threshold: {analysis_result['aggregation_value']}" + ) + return None + + # Create aggregated knowledge + team_knowledge_id = await self.create_team_knowledge( + team_id=team_id, + title=analysis_result["title"], + content=analysis_result["content"], + content_type=analysis_result["content_type"], + knowledge_category=analysis_result["category"], + source_type=SourceType.AGENT_CONTRIBUTION, + contributing_agents=[agent_id], + aggregation_method=aggregation_method, + team_relevance_score=analysis_result["aggregation_value"], + metadata=analysis_result["metadata"], + tags=analysis_result["tags"], + ) + + # Mark original memories as aggregated + await conn.execute( + """ + UPDATE agent_memory + SET propagated_to_team = TRUE, team_context_id = $2 + WHERE id = ANY($1) + """, + agent_memory_ids, + team_knowledge_id, + ) + + self.aggregations_performed += 1 + + logger.info( + f"Aggregated {len(agent_memory_ids)} agent memories into team knowledge {team_knowledge_id}" + ) + return team_knowledge_id + + async def get_team_knowledge_context( + self, + team_id: str, + task_context: Dict[str, Any], + agent_id: Optional[str] = None, + max_context_items: int = 5, + ) -> Dict[str, Any]: + """Get relevant team knowledge for task execution context""" + + # Build context query from task information + context_query = self._build_context_query(task_context) + + # Search for relevant knowledge + search_results = await self.search_team_knowledge( + team_id=team_id, + query=context_query, + limit=max_context_items, + min_similarity=0.4, + ) + + # Get team statistics + team_stats = await self.get_team_knowledge_stats(team_id) + + # Build context + context = { + "team_id": team_id, + "relevant_knowledge": [ + { + "id": result.team_knowledge.id, + "title": result.team_knowledge.title, + "content": ( + result.team_knowledge.content[:500] + "..." + if len(result.team_knowledge.content) > 500 + else result.team_knowledge.content + ), + "category": result.team_knowledge.knowledge_category.value, + "relevance_score": result.combined_score, + "usage_stats": { + "adoption_rate": result.team_knowledge.agent_adoption_rate, + "effectiveness": result.team_knowledge.effectiveness_score, + }, + } + for result in search_results + ], + "team_knowledge_stats": team_stats, + "context_query": context_query, + "generated_at": datetime.now().isoformat(), + } + + return context + + async def update_knowledge_effectiveness( + self, + team_knowledge_id: str, + agent_id: str, + task_success: bool, + feedback_score: Optional[float] = None, + usage_context: Optional[Dict[str, Any]] = None, + ): + """Update knowledge effectiveness based on agent usage""" + + async with self.pool.acquire() as conn: + # Get current knowledge + knowledge = await conn.fetchrow( + """ + SELECT * FROM team_knowledge_base WHERE id = $1 + """, + team_knowledge_id, + ) + + if not knowledge: + return + + # Calculate new effectiveness score + success_weight = 1.0 if task_success else -0.3 + feedback_weight = (feedback_score or 0.5) - 0.5 + + # Update effectiveness with exponential moving average + current_effectiveness = knowledge["effectiveness_score"] + new_effectiveness = ( + current_effectiveness * 0.8 + (success_weight + feedback_weight) * 0.2 + ) + new_effectiveness = max(0.0, min(1.0, new_effectiveness)) + + # Update agent adoption tracking + contributing_agents = knowledge["contributing_agents"] or [] + if agent_id not in contributing_agents: + contributing_agents.append(agent_id) + + # Calculate adoption rate (agents who used it / total team agents) + team_agent_count = await conn.fetchval( + """ + SELECT COUNT(*) FROM agents WHERE team_id = $1 + """, + knowledge["team_id"], + ) + + adoption_rate = len(contributing_agents) / max(1, team_agent_count) + + # Update knowledge + await conn.execute( + """ + UPDATE team_knowledge_base + SET effectiveness_score = $2, + agent_adoption_rate = $3, + contributing_agents = $4, + last_accessed = NOW(), + updated_at = NOW() + WHERE id = $1 + """, + team_knowledge_id, + new_effectiveness, + adoption_rate, + contributing_agents, + ) + + logger.debug( + f"Updated knowledge {team_knowledge_id} effectiveness: {new_effectiveness:.2f}, adoption: {adoption_rate:.2f}" + ) + + async def get_team_knowledge_stats(self, team_id: str) -> Dict[str, Any]: + """Get comprehensive team knowledge statistics""" + + async with self.pool.acquire() as conn: + # Basic statistics + basic_stats = await conn.fetchrow( + """ + SELECT + COUNT(*) as total_knowledge, + COUNT(DISTINCT knowledge_category) as categories, + COUNT(DISTINCT unnest(contributing_agents)) as contributing_agents, + AVG(team_relevance_score) as avg_relevance, + AVG(effectiveness_score) as avg_effectiveness, + AVG(agent_adoption_rate) as avg_adoption_rate + FROM team_knowledge_base + WHERE team_id = $1 + """, + team_id, + ) + + # Category breakdown + category_stats = await conn.fetch( + """ + SELECT + knowledge_category, + COUNT(*) as count, + AVG(effectiveness_score) as avg_effectiveness, + AVG(agent_adoption_rate) as avg_adoption + FROM team_knowledge_base + WHERE team_id = $1 + GROUP BY knowledge_category + ORDER BY count DESC + """, + team_id, + ) + + # Most effective knowledge + top_knowledge = await conn.fetch( + """ + SELECT + title, + knowledge_category, + effectiveness_score, + agent_adoption_rate + FROM team_knowledge_base + WHERE team_id = $1 + ORDER BY effectiveness_score DESC + LIMIT 5 + """, + team_id, + ) + + return { + "team_id": team_id, + "basic_stats": dict(basic_stats) if basic_stats else {}, + "category_breakdown": [dict(cat) for cat in category_stats], + "top_knowledge": [dict(know) for know in top_knowledge], + "generated_at": datetime.now().isoformat(), + } + + async def _search_team_specific_knowledge( + self, + conn, + team_id: str, + query_embedding: List[float], + categories: Optional[List[KnowledgeCategory]], + content_types: Optional[List[ContentType]], + limit: int, + min_similarity: float, + ) -> List[TeamKnowledgeSearchResult]: + """Search team-specific knowledge base""" + + # Build query conditions + where_conditions = ["team_id = $2"] + params = [query_embedding, team_id] + param_idx = 3 + + if categories: + where_conditions.append(f"knowledge_category = ANY(${param_idx})") + params.append([cat.value for cat in categories]) + param_idx += 1 + + if content_types: + where_conditions.append(f"content_type = ANY(${param_idx})") + params.append([ct.value for ct in content_types]) + param_idx += 1 + + where_conditions.append(f"(1 - (embedding <=> $1)) >= ${param_idx}") + params.append(min_similarity) + param_idx += 1 + + where_clause = "WHERE " + " AND ".join(where_conditions) + + results = await conn.fetch( + f""" + SELECT + *, + (1 - (embedding <=> $1)) as similarity_score + FROM team_knowledge_base + {where_clause} + ORDER BY similarity_score DESC, effectiveness_score DESC + LIMIT ${param_idx} + """, # nosec B608 -- where clause is fixed fragments with $N placeholders; all values bound as query params + *params, + limit, + ) + + search_results = [] + for row in results: + team_knowledge = self._row_to_team_knowledge(row) + + # Calculate team fit score based on adoption and effectiveness + team_fit_score = ( + team_knowledge.agent_adoption_rate * 0.4 + + team_knowledge.effectiveness_score * 0.6 + ) + + combined_score = ( + float(row["similarity_score"]) * 0.4 + + team_knowledge.team_relevance_score * 0.3 + + team_fit_score * 0.3 + ) + + search_results.append( + TeamKnowledgeSearchResult( + team_knowledge=team_knowledge, + similarity_score=float(row["similarity_score"]), + relevance_score=team_knowledge.team_relevance_score, + team_fit_score=team_fit_score, + combined_score=combined_score, + ) + ) + + return search_results + + async def _search_org_knowledge_for_team( + self, + conn, + team_id: str, + query: str, + categories: Optional[List[KnowledgeCategory]], + content_types: Optional[List[ContentType]], + limit: int, + min_similarity: float, + ) -> List[TeamKnowledgeSearchResult]: + """Search organization knowledge filtered for team relevance""" + + # Get organization ID for the team + org_id = await conn.fetchval( + """ + SELECT organization_id FROM teams WHERE id = $1 + """, + team_id, + ) + + if not org_id: + return [] + + # Search organization knowledge + org_results = await self.org_rag_manager.search_knowledge( + organization_id=str(org_id), + query=query, + categories=categories, + content_types=content_types, + limit=limit * 2, # Get more to filter for team relevance + min_similarity=min_similarity, + requester_team_id=team_id, + ) + + # Convert to team knowledge search results with team relevance scoring + team_results = [] + for org_result in org_results: + # Calculate team relevance based on source and usage + team_relevance = await self._calculate_team_relevance( + conn, team_id, org_result.knowledge + ) + + if team_relevance >= self.min_team_relevance: + # Create pseudo team knowledge for consistent interface + pseudo_team_knowledge = TeamKnowledge( + id=org_result.knowledge.id, + team_id=team_id, + organization_id=org_result.knowledge.organization_id, + title=org_result.knowledge.title, + content=org_result.knowledge.content, + content_type=org_result.knowledge.content_type, + knowledge_category=org_result.knowledge.knowledge_category, + embedding=org_result.knowledge.embedding, + source_type=org_result.knowledge.source_type, + contributing_agents=[], + source_knowledge_ids=[org_result.knowledge.id], + aggregation_method="organization_filter", + team_relevance_score=team_relevance, + agent_adoption_rate=0.0, + effectiveness_score=org_result.knowledge.success_correlation, + visibility_level=org_result.knowledge.visibility_level, + metadata=org_result.knowledge.metadata, + tags=org_result.knowledge.tags, + created_at=org_result.knowledge.created_at, + updated_at=org_result.knowledge.updated_at, + last_accessed=org_result.knowledge.last_accessed, + ) + + combined_score = ( + org_result.similarity_score * 0.5 + team_relevance * 0.5 + ) + + team_results.append( + TeamKnowledgeSearchResult( + team_knowledge=pseudo_team_knowledge, + similarity_score=org_result.similarity_score, + relevance_score=org_result.relevance_score, + team_fit_score=team_relevance, + combined_score=combined_score, + ) + ) + + return team_results[:limit] + + async def _calculate_team_relevance( + self, conn, team_id: str, org_knowledge + ) -> float: + """Calculate how relevant organization knowledge is for a specific team""" + + relevance_factors = [] + + # Factor 1: Source team match + if org_knowledge.source_team_id == team_id: + relevance_factors.append(1.0) + elif org_knowledge.source_team_id: + # Check if source team is similar to current team + team_similarity = await self._calculate_team_similarity( + conn, team_id, org_knowledge.source_team_id + ) + relevance_factors.append(team_similarity) + else: + relevance_factors.append(0.3) # No team context + + # Factor 2: Category relevance to team's work + team_categories = await self._get_team_primary_categories(conn, team_id) + if org_knowledge.knowledge_category.value in team_categories: + relevance_factors.append(0.9) + else: + relevance_factors.append(0.4) + + # Factor 3: Usage by team agents + team_usage = ( + await conn.fetchval( + """ + SELECT COUNT(DISTINCT source_agent_id)::float / NULLIF( + (SELECT COUNT(*) FROM agents WHERE team_id = $1), 0 + ) + FROM organization_knowledge_base + WHERE id = $2 AND source_agent_id IN ( + SELECT id FROM agents WHERE team_id = $1 + ) + """, + team_id, + org_knowledge.id, + ) + or 0.0 + ) + relevance_factors.append(team_usage) + + # Factor 4: Base quality and relevance + relevance_factors.append(org_knowledge.quality_score) + relevance_factors.append(org_knowledge.relevance_score) + + # Calculate weighted average + weights = [0.3, 0.25, 0.25, 0.1, 0.1] + team_relevance = sum(f * w for f, w in zip(relevance_factors, weights)) + + return min(1.0, max(0.0, team_relevance)) + + async def _calculate_team_similarity( + self, conn, team_id1: str, team_id2: str + ) -> float: + """Calculate similarity between two teams based on their work patterns""" + + # Simple implementation based on team type and settings + team_info = await conn.fetch( + """ + SELECT id, team_type, settings FROM teams + WHERE id IN ($1, $2) + """, + team_id1, + team_id2, + ) + + if len(team_info) != 2: + return 0.0 + + team1, team2 = team_info + + # Type similarity + type_similarity = 1.0 if team1["team_type"] == team2["team_type"] else 0.5 + + # Settings similarity (simplified) + settings1 = team1["settings"] or {} + settings2 = team2["settings"] or {} + + common_keys = set(settings1.keys()) & set(settings2.keys()) + if common_keys: + settings_similarity = sum( + 1.0 if settings1.get(key) == settings2.get(key) else 0.0 + for key in common_keys + ) / len(common_keys) + else: + settings_similarity = 0.5 + + return type_similarity * 0.7 + settings_similarity * 0.3 + + async def _get_team_primary_categories(self, conn, team_id: str) -> List[str]: + """Get primary knowledge categories this team works with""" + + categories = await conn.fetch( + """ + SELECT knowledge_category, COUNT(*) as usage_count + FROM team_knowledge_base + WHERE team_id = $1 + GROUP BY knowledge_category + ORDER BY usage_count DESC + LIMIT 3 + """, + team_id, + ) + + return [cat["knowledge_category"] for cat in categories] + + def _generate_embedding(self, text: str) -> List[float]: + """Generate embedding for text using sentence transformers""" + try: + embedding = self.embedding_model.encode(text, convert_to_tensor=False) + return embedding.tolist() + except Exception as e: + logger.error(f"Error generating embedding: {e}") + return [0.0] * self.embedding_dim + + def _row_to_team_knowledge(self, row) -> TeamKnowledge: + """Convert database row to TeamKnowledge object""" + return TeamKnowledge( + id=str(row["id"]), + team_id=str(row["team_id"]), + organization_id=str(row["organization_id"]), + title=row["title"], + content=row["content"], + content_type=ContentType(row["content_type"]), + knowledge_category=KnowledgeCategory(row["knowledge_category"]), + embedding=row["embedding"] if row["embedding"] else None, + source_type=SourceType(row["source_type"]), + contributing_agents=row["contributing_agents"] or [], + source_knowledge_ids=row["source_knowledge_ids"] or [], + aggregation_method=row["aggregation_method"], + team_relevance_score=row["team_relevance_score"], + agent_adoption_rate=row["agent_adoption_rate"], + effectiveness_score=row["effectiveness_score"], + visibility_level=VisibilityLevel(row["visibility_level"]), + metadata=( + json.loads(row["metadata"]) + if isinstance(row["metadata"], str) + else row["metadata"] + ), + tags=row["tags"] or [], + created_at=row["created_at"], + updated_at=row["updated_at"], + last_accessed=row["last_accessed"], + ) + + def _build_context_query(self, task_context: Dict[str, Any]) -> str: + """Build a search query from task context""" + query_parts = [] + + if task_context.get("task_type"): + query_parts.append(task_context["task_type"]) + + if task_context.get("description"): + query_parts.append(task_context["description"]) + + if task_context.get("technologies"): + query_parts.extend(task_context["technologies"]) + + if task_context.get("domain"): + query_parts.append(task_context["domain"]) + + return " ".join(query_parts) + + async def _analyze_memories_for_aggregation( + self, agent_memories: List + ) -> Dict[str, Any]: + """Analyze agent memories to determine if they should be aggregated""" + + if not agent_memories: + return {"aggregation_value": 0.0} + + # Simple aggregation analysis + # In practice, this could use more sophisticated NLP + + # Calculate average confidence and success correlation + avg_confidence = sum(mem["confidence_score"] for mem in agent_memories) / len( + agent_memories + ) + avg_success = sum( + mem.get("success_correlation", 0.0) for mem in agent_memories + ) / len(agent_memories) + + # Find common themes + all_content = " ".join(mem["content"] for mem in agent_memories) + + # Determine primary category + categories = [mem.get("memory_type", "general") for mem in agent_memories] + primary_category = ( + max(set(categories), key=categories.count) if categories else "general" + ) + + # Create aggregated content (simplified) + title = f"Team Knowledge: {primary_category.replace('_', ' ').title()}" + content = ( + f"Aggregated knowledge from {len(agent_memories)} agent experiences:\n\n" + + all_content[:1000] + ) + + # Map memory type to knowledge category + category_mapping = { + "code_pattern": KnowledgeCategory.DEVELOPMENT, + "task_outcome": KnowledgeCategory.PROCESS, + "debugging": KnowledgeCategory.TROUBLESHOOTING, + "optimization": KnowledgeCategory.DEVELOPMENT, + "testing": KnowledgeCategory.TESTING, + } + + knowledge_category = category_mapping.get( + primary_category, KnowledgeCategory.DEVELOPMENT + ) + + # Determine content type + content_type = ( + ContentType.CODE if "code" in primary_category else ContentType.TEXT + ) + + # Calculate aggregation value + aggregation_value = min(1.0, (avg_confidence + avg_success) / 2.0) + + return { + "aggregation_value": aggregation_value, + "title": title, + "content": content, + "content_type": content_type, + "category": knowledge_category, + "metadata": { + "source_memory_count": len(agent_memories), + "avg_confidence": avg_confidence, + "avg_success_correlation": avg_success, + "primary_type": primary_category, + }, + "tags": [primary_category, "aggregated", "agent_contribution"], + }