Spaces:
Running
Running
File size: 3,146 Bytes
a55ea57 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | """Shared semantic agent-memory operations over existing BrainAPI handlers."""
from __future__ import annotations
from typing import Any
from src.services.api.constants.requests import GetContextRequestBody, IngestionRequestBody
from src.constants.tasks.ingestion import IngestionTaskTextArgs
from src.services.api.routes.ingest import ingest_data
from src.services.api.routes.retrieve import get_context, retrieve
from src.services.api.routes.tasks import get_task
from src.services.api.controllers.entities import get_entity_info
from src.services.api.controllers.retrieve import retrieve_neighbors
def _json(value: Any) -> Any:
if hasattr(value, "model_dump"):
return value.model_dump(mode="json")
return value
async def memory_search(query: str, limit: int, brain_id: str) -> list[dict]:
response = await retrieve(query, limit, "", brain_id)
results: list[dict] = []
for kind, values in (
("chunk", response.data),
("observation", response.observations),
("relationship", response.relationships),
):
for value in values:
item = _json(value)
results.append(
{
"id": str(item.get("id") or item.get("uuid") or ""),
"type": kind,
"text": item.get("text") or item.get("description") or "",
"score": item.get("score"),
"metadata": item.get("metadata") or {},
"data": item,
}
)
return results[:limit]
async def memory_context(query: str, max_tokens: int, brain_id: str) -> dict:
# Existing context retrieval is fact/passage bounded rather than token bounded.
# Convert the stable agent hint into conservative existing limits.
max_facts = max(1, min(100, max_tokens // 100))
request = GetContextRequestBody(text=query, max_facts=max_facts)
response = await get_context(request, brain_id)
payload = _json(response)
return {
"context": payload.get("text_context", ""),
"sources": payload.get("source_passages", []),
"triples": payload.get("triples", []),
}
async def memory_store(
*,
text: str,
metadata: dict,
brain_id: str,
request,
) -> dict:
body = IngestionRequestBody(
data=IngestionTaskTextArgs(text_data=text),
meta_keys=metadata or None,
)
response = await ingest_data(body, request, brain_id)
payload = response.body.decode("utf-8")
import json
return json.loads(payload)
async def memory_task_status(task_id: str, brain_id: str) -> dict:
return await get_task(task_id, brain_id)
async def memory_neighbors(entity_uuid: str, limit: int, brain_id: str) -> Any:
"""Reuse canonical graph-neighbor retrieval for semantic MCP access."""
return await retrieve_neighbors(uuid=entity_uuid, limit=limit, brain_id=brain_id)
async def memory_entity(target: str, max_depth: int, brain_id: str) -> Any:
"""Reuse canonical entity information retrieval for semantic MCP access."""
return await get_entity_info(target, target, max_depth, brain_id)
|