diff --git a/src/brain.py b/src/brain.py index bc5e21d..876fb98 100644 --- a/src/brain.py +++ b/src/brain.py @@ -1,7 +1,13 @@ """Core brain operations — capture, search, recent, stats.""" + import json + from src import db, embeddings -from src.models import CaptureRequest, Thought, SearchRequest, SearchResult, Stats +from src.models import CaptureRequest, SearchRequest, SearchResult, Stats, Thought + + +def _meta(raw) -> dict: + return json.loads(raw) if isinstance(raw, str) else dict(raw) async def capture(req: CaptureRequest) -> Thought: @@ -15,25 +21,32 @@ async def capture(req: CaptureRequest) -> Thought: """ INSERT INTO thoughts (content, embedding, source, metadata) VALUES ($1, $2::vector, $3, $4::jsonb) - RETURNING id::text, content, source, metadata, created_at, embedding IS NOT NULL AS embedded + RETURNING id::text, content, source, metadata, created_at, + embedding IS NOT NULL AS embedded """, - req.content, vec, req.source, json.dumps(req.metadata), + req.content, + vec, + req.source, + json.dumps(req.metadata), ) else: row = await conn.fetchrow( """ INSERT INTO thoughts (content, source, metadata) VALUES ($1, $2, $3::jsonb) - RETURNING id::text, content, source, metadata, created_at, embedding IS NOT NULL AS embedded + RETURNING id::text, content, source, metadata, created_at, + embedding IS NOT NULL AS embedded """, - req.content, req.source, json.dumps(req.metadata), + req.content, + req.source, + json.dumps(req.metadata), ) return Thought( id=row["id"], content=row["content"], source=row["source"], - metadata=json.loads(row["metadata"]) if isinstance(row["metadata"], str) else row["metadata"], + metadata=_meta(row["metadata"]), created_at=row["created_at"], embedded=row["embedded"], ) @@ -58,7 +71,9 @@ async def search(req: SearchRequest) -> list[SearchResult]: ORDER BY embedding <=> $1::vector LIMIT $3 """, - vec, req.source, req.limit, + vec, + req.source, + req.limit, ) else: rows = await conn.fetch( @@ -70,7 +85,8 @@ async def search(req: SearchRequest) -> list[SearchResult]: ORDER BY embedding <=> $1::vector LIMIT $2 """, - vec, req.limit, + vec, + req.limit, ) return [ @@ -80,7 +96,7 @@ async def search(req: SearchRequest) -> list[SearchResult]: source=r["source"], similarity=float(r["similarity"]), created_at=r["created_at"], - metadata=json.loads(r["metadata"]) if isinstance(r["metadata"], str) else r["metadata"], + metadata=_meta(r["metadata"]), ) for r in rows ] diff --git a/src/db.py b/src/db.py index 9501d75..10bcaec 100644 --- a/src/db.py +++ b/src/db.py @@ -1,4 +1,5 @@ import asyncpg + from src.config import settings _pool: asyncpg.Pool | None = None diff --git a/src/embeddings.py b/src/embeddings.py index 742c95b..48a7ce4 100644 --- a/src/embeddings.py +++ b/src/embeddings.py @@ -1,4 +1,5 @@ import httpx + from src.config import settings diff --git a/src/main.py b/src/main.py index ee31703..dcb67ed 100644 --- a/src/main.py +++ b/src/main.py @@ -1,4 +1,5 @@ """OpenBrain — FastAPI REST + MCP server (single process).""" + import contextlib import logging @@ -40,7 +41,8 @@ async def capture(content: str, source: str = "unknown", metadata: dict | None = source: Which agent or tool is capturing this (e.g. 'claude-code', 'codex') metadata: Optional key/value pairs (tags, project, etc.) """ - thought = await brain.capture(CaptureRequest(content=content, source=source, metadata=metadata or {})) + req = CaptureRequest(content=content, source=source, metadata=metadata or {}) + thought = await brain.capture(req) return thought.model_dump(mode="json") diff --git a/src/models.py b/src/models.py index 14fc259..d8deeb6 100644 --- a/src/models.py +++ b/src/models.py @@ -1,5 +1,6 @@ from datetime import datetime from typing import Any + from pydantic import BaseModel