added improvements

Author: marcorasera54

src/services/api/routes/retrieve.py

@@ -159,51 +159,57 @@ async def get_entities(
     )
 
 @retrieve_router.get(path="/structured-data/types")
-async def get_structured_data_types():
+async def get_structured_data_types(
+    brain_id: str = "default",
+):
     """
     Get all unique types from structured data.
     """
-    return await get_structured_data_types_controller()
+    return await get_structured_data_types_controller(brain_id)
 
 @retrieve_router.get(path="/structured-data/{id}")
 async def get_structured_data_by_id(
     id: str,
+    brain_id: str = "default",
 ):
     """
     Get structured data by ID.
     """
-    return await get_structured_data_by_id_controller(id)
+    return await get_structured_data_by_id_controller(id, brain_id)
 
 @retrieve_router.get(path="/structured-data")
 async def get_structured_data_list(
     limit: int = 10,
     skip: int = 0,
     types: Optional[str] = None,
     query_text: Optional[str] = None,
+    brain_id: str = "default",
 ):
     """
     Get a list of structured data.
     """
     if types:
         types = types.split(",")
-    return await get_structured_data_list_controller(limit, skip, types, query_text)
+    return await get_structured_data_list_controller(limit, skip, types, query_text, brain_id)
 
 @retrieve_router.get(path="/observations/labels")
-async def get_observation_labels():
+async def get_observation_labels(
+    brain_id: str = "default",
+):
     """
     Get all unique labels from observations.
     """
-    return await get_observation_labels_controller()
+    return await get_observation_labels_controller(brain_id)
 
 @retrieve_router.get(path="/observations/{id}")
 async def get_observation_by_id(
     id: str,
+    brain_id: str = "default",
 ):
     """
     Get observation by ID.
     """
-    return await get_observation_by_id_controller(id)
-
+    return await get_observation_by_id_controller(id, brain_id)
 
 @retrieve_router.get(path="/observations")
 async def get_observations_list(
@@ -212,10 +218,11 @@ async def get_observations_list(
     resource_id: Optional[str] = None,
     labels: Optional[str] = None,
     query_text: Optional[str] = None,
+    brain_id: str = "default",
 ):
     """
     Get a list of observations.
     """
     if labels:
         labels = labels.split(",")
-    return await get_observations_list_controller(limit, skip, resource_id, labels, query_text)
+    return await get_observations_list_controller(limit, skip, resource_id, labels, query_text, brain_id)

src/services/api/routes/tasks.py

@@ -9,37 +9,40 @@
 """
 import json
 from src.utils.logging import log
-from fastapi import APIRouter, HTTPException
+from fastapi import APIRouter
 from src.services.kg_agent.main import cache_adapter
 
 tasks_router = APIRouter(prefix="/tasks", tags=["tasks"])
 
-
 @tasks_router.get("/{task_id}")
 async def get_task(task_id: str, brain_id: str = "default"):
     """
     Get the result of a task by its ID.
     """
     try:
         str_result = cache_adapter.get(f"task:{task_id}", brain_id=brain_id)
-
         if str_result is None:
-            raise HTTPException(status_code=404, detail="Task not found")
-
+            return {
+                "task_id": task_id,
+                "status": "pending",
+                "result": {"message": "Task is still processing or not found"}
+            }
         if isinstance(str_result, bytes):
             result = json.loads(str_result.decode("utf-8"))
         else:
             result = json.loads(str_result)
-
         return {
             "task_id": task_id,
-            "status": result.get("status", "pending"),
+            "status": result.get("status", "unknown"),
             "result": result
         }
     except Exception as e:
         log(f"Error in get_task: {type(e).__name__}: {str(e)}")
         return {
             "task_id": task_id,
             "status": "error",
-            "result": {"error": str(e)}
+            "result": {
+                "error": str(e),
+                "error_type": type(e).__name__
+            }
         }