devbrain_service.py

import uvicorn
from fastapi import FastAPI, HTTPException, BackgroundTasks
from pydantic import BaseModel
from typing import Optional, Dict, Any
import devbrain_core
import asyncio
import get_dashboard_data

app = FastAPI(title="DevBrain Core Service")

# Track active background tasks to report to the frontend
active_tasks = {
    "capture": False,
    "optimize": False,
    "last_capture_time": None,
    "last_optimize_time": None
}

# Initialize memory connection when the service starts
@app.on_event("startup")
async def startup_event():
    print("[DevBrain Service] Initializing Cognee local configurations...")
    await devbrain_core.init_memory()

class CaptureRequest(BaseModel):
    type: str = "manual_note"
    content: str
    metadata: Optional[Dict[str, Any]] = {}

class RecallRequest(BaseModel):
    query: str

class PurgeRequest(BaseModel):
    dataset: str

async def run_capture_in_background(event_type: str, content: str, metadata: dict):
    active_tasks["capture"] = True
    start_time = asyncio.get_event_loop().time()
    try:
        await devbrain_core.capture_memory_event(event_type, content, metadata)
    except Exception as e:
        print(f"[DevBrain Service] Background capture error: {e}")
    finally:
        active_tasks["capture"] = False
        active_tasks["last_capture_time"] = asyncio.get_event_loop().time() - start_time

async def run_optimize_in_background():
    active_tasks["optimize"] = True
    start_time = asyncio.get_event_loop().time()
    try:
        await devbrain_core.optimize_memory_result()
    except Exception as e:
        print(f"[DevBrain Service] Background optimize error: {e}")
    finally:
        active_tasks["optimize"] = False
        active_tasks["last_optimize_time"] = asyncio.get_event_loop().time() - start_time

@app.get("/health")
async def health():
    h = devbrain_core.get_memory_health()
    h["active_tasks"] = active_tasks
    return {"ok": True, "data": h}

@app.get("/datasets")
async def datasets():
    return {"ok": True, "data": devbrain_core.list_datasets()}

@app.get("/dashboard_data")
async def dashboard_data():
    data = await get_dashboard_data.get_dashboard_data()
    # Inject active tasks for frontend stopwatch syncing
    data["active_tasks"] = active_tasks
    return data

@app.post("/optimize")
async def optimize(background_tasks: BackgroundTasks):
    if active_tasks["optimize"]:
        return {"ok": False, "error": "An optimization task is already in progress."}
    background_tasks.add_task(run_optimize_in_background)
    return {"ok": True, "status": "queued", "message": "Graph optimization queued in the background."}

@app.post("/purge")
async def purge(req: PurgeRequest):
    result = await devbrain_core.purge_memory_result(req.dataset)
    return result

@app.post("/capture")
async def capture(req: CaptureRequest, background_tasks: BackgroundTasks):
    if active_tasks["capture"]:
        return {"ok": False, "error": "A capture task is already in progress."}
    background_tasks.add_task(run_capture_in_background, req.type, req.content, req.metadata)
    return {"ok": True, "status": "queued", "source": "cognee"}

@app.post("/recall")
async def recall(req: RecallRequest):
    result = await devbrain_core.query_memory_result(req.query)
    return result

def run_service(host: str = "127.0.0.1", port: int = 8000):
    uvicorn.run("devbrain_service:app", host=host, port=port, log_level="info")

if __name__ == "__main__":
    run_service()