changed types

Author: MayaCrmi

multi-agent/elements/retrievers/docs_dataflow/config.py

@@ -1,4 +1,4 @@
-from typing import Literal, Optional, List, Union
+from typing import Dict, List, Literal, Optional
 from pydantic import Field
 from elements.retrievers.common.base_config import BaseRetrieverConfig
 from core.field_hints import ActionHint, HintType, SelectionType
@@ -31,7 +31,7 @@ class DocsDataflowRetrieverConfig(BaseRetrieverConfig):
         description="Request timeout in seconds"
     )
 
-    docs: Optional[List[Union[dict, str]]] = Field(
+    docs: Optional[List[Dict]] = Field(
         default=None,
         description="Filter results to specific documents",
         json_schema_extra=ActionHint(

multi-agent/elements/retrievers/docs_dataflow/docs_dataflow_retriever.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union
+from typing import Dict, List, Optional
 from elements.retrievers.common.base_retriever import BaseRetriever
 from elements.providers.dataflow_client.config import DataflowProviderConfig
 from elements.providers.dataflow_client.dataflow_provider_factory import DataflowProviderFactory
@@ -15,7 +15,7 @@ def __init__(
             top_k_results: int,
             threshold: float,
             timeout: float = 30.0,
-            docs: Optional[List[Union[dict, str]]] = None,
+            docs: Optional[List[Dict]] = None,
             tags: Optional[List[str]] = None,
     ):
         self.threshold = threshold
@@ -31,8 +31,8 @@ def __init__(
     def retrieve(self, query: str) -> List[dict]:
         context = get_current_context()
 
-        # Extract document IDs from docs list (handles both dict and string items)
-        doc_ids = [doc.get('id') if isinstance(doc, dict) else doc for doc in (self.docs or [])] or None
+        # Extract document IDs from docs list
+        doc_ids = [doc['id'] for doc in self.docs] if self.docs else None
 
         response = self._provider.query(
             query=query,