hedit 0.7.5a2__tar.gz → 0.7.5.dev2__tar.gz

{hedit-0.7.5a2/hedit.egg-info → hedit-0.7.5.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hedit
-Version: 0.7.5a2
+Version: 0.7.5.dev2
 Summary: Multi-agent system for HED annotation generation and validation
 Author-email: Annotation Garden Initiative <info@annotation.garden>
 License-Expression: MIT

{hedit-0.7.5a2 → hedit-0.7.5.dev2/hedit.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hedit
-Version: 0.7.5a2
+Version: 0.7.5.dev2
 Summary: Multi-agent system for HED annotation generation and validation
 Author-email: Annotation Garden Initiative <info@annotation.garden>
 License-Expression: MIT

{hedit-0.7.5a2 → hedit-0.7.5.dev2}/hedit.egg-info/SOURCES.txt

@@ -66,7 +66,9 @@ tests/test_github_client.py
 tests/test_hed_lsp.py
 tests/test_integration_openrouter.py
 tests/test_json_schema_loader.py
+tests/test_keyword_extraction.py
 tests/test_litellm_llm.py
+tests/test_no_extend_propagation.py
 tests/test_openrouter_llm.py
 tests/test_schema_loader.py
 tests/test_security.py

{hedit-0.7.5a2 → hedit-0.7.5.dev2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "hedit"
-version = "0.7.5a2"
+version = "0.7.5.dev2"
 description = "Multi-agent system for HED annotation generation and validation"
 readme = "PKG_README.md"
 requires-python = ">=3.12"

{hedit-0.7.5a2 → hedit-0.7.5.dev2}/src/agents/workflow.py

@@ -9,6 +9,7 @@ import time
 from pathlib import Path
 
 from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import HumanMessage, SystemMessage
 from langgraph.graph import END, StateGraph
 
 from src.agents.annotation_agent import AnnotationAgent
@@ -17,6 +18,7 @@ from src.agents.evaluation_agent import EvaluationAgent
 from src.agents.feedback_summarizer import FeedbackSummarizer
 from src.agents.state import HedAnnotationState
 from src.agents.validation_agent import ValidationAgent
+from src.utils import extract_text_content
 from src.utils.schema_loader import HedSchemaLoader
 from src.validation.hed_lsp import HedLspClient, is_hed_lsp_available
 
@@ -63,8 +65,8 @@ class HedAnnotationWorkflow:
         """
         # Store schema directory (None means use HED library to fetch from GitHub)
         self.schema_dir = schema_dir
-        # Enable semantic search only if hed-lsp CLI is available
-        self.enable_semantic_search = enable_semantic_search and is_hed_lsp_available()
+        # Keyword extraction always runs; LSP enrichment requires hed-lsp CLI
+        self.enable_semantic_search = enable_semantic_search
 
         # Initialize legacy schema loader for validation
         self.schema_loader = HedSchemaLoader()
@@ -74,6 +76,9 @@ class HedAnnotationWorkflow:
         assess_llm = assessment_llm or llm
         feed_llm = feedback_llm or llm
 
+        # Store feedback LLM for keyword extraction (cheap/fast model)
+        self.feedback_llm = feed_llm
+
         # Initialize agents with JSON schema support and per-agent LLMs
         self.annotation_agent = AnnotationAgent(llm, schema_dir=self.schema_dir)
         self.validation_agent = ValidationAgent(
@@ -85,15 +90,19 @@ class HedAnnotationWorkflow:
         self.assessment_agent = AssessmentAgent(assess_llm, schema_dir=self.schema_dir)
         self.feedback_summarizer = FeedbackSummarizer(feed_llm)
 
-        # Initialize hed-lsp client for semantic search
+        # Initialize hed-lsp client for semantic search (optional enrichment)
         self.hed_lsp_client: HedLspClient | None = None
-        if self.enable_semantic_search:
+        if self.enable_semantic_search and is_hed_lsp_available():
             try:
                 self.hed_lsp_client = HedLspClient()
                 logger.info("[WORKFLOW] hed-lsp CLI available for semantic tag suggestions")
             except RuntimeError as e:
                 logger.warning(f"[WORKFLOW] hed-lsp CLI not available: {e}")
-                self.enable_semantic_search = False
+        elif self.enable_semantic_search:
+            logger.info(
+                "[WORKFLOW] hed-lsp CLI not in PATH; "
+                "keyword extraction will run without LSP enrichment"
+            )
 
         # Build graph
         self.graph = self._build_graph()
@@ -156,11 +165,58 @@ class HedAnnotationWorkflow:
 
         return workflow.compile()  # type: ignore[return-value]
 
+    async def _extract_keywords(self, description: str) -> list[str]:
+        """Extract HED-relevant keywords from a natural language description.
+
+        Uses the feedback LLM (cheap/fast model) to identify key concepts
+        that can be mapped to HED tags via the LSP suggest tool.
+
+        Args:
+            description: Natural language event or image description
+
+        Returns:
+            List of extracted keywords (max 20)
+        """
+        system_prompt = (
+            "You are a keyword extractor for neuroscience event descriptions. "
+            "Extract the most important concepts that could map to HED "
+            "(Hierarchical Event Descriptors) tags.\n\n"
+            "Extract:\n"
+            "- Objects/entities (person, car, button, screen, face, etc.)\n"
+            "- Actions/events (pressing, flashing, appearing, moving, etc.)\n"
+            "- Properties/attributes (red, large, fast, loud, etc.)\n"
+            "- Spatial relationships (left, center, above, etc.)\n"
+            "- Temporal aspects (onset, offset, duration, etc.)\n"
+            "- Sensory modalities (visual, auditory, tactile, etc.)\n\n"
+            "Return ONLY a comma-separated list of single words or short phrases "
+            "(2-3 words max). Return at most 20 keywords. "
+            "Do not include any other text, explanation, or formatting."
+        )
+
+        try:
+            response = await self.feedback_llm.ainvoke(
+                [
+                    SystemMessage(content=system_prompt),
+                    HumanMessage(content=f"Description: {description}"),
+                ]
+            )
+            raw_text = extract_text_content(response.content)
+            # Parse comma-separated keywords, strip whitespace, filter empty
+            keywords = [kw.strip() for kw in raw_text.split(",") if kw.strip()]
+            # Limit to 20 keywords
+            keywords = keywords[:20]
+            logger.info(f"[WORKFLOW] Extracted {len(keywords)} keywords: {keywords}")
+            return keywords
+        except Exception as e:
+            logger.warning("[WORKFLOW] Keyword extraction failed: %s", e, exc_info=True)
+            return []
+
     async def _semantic_preprocess_node(self, state: HedAnnotationState) -> dict:
-        """Semantic preprocessing node: Use hed-lsp CLI to suggest relevant tags.
+        """Semantic preprocessing node: Extract keywords and suggest HED tags.
 
         This node runs before annotation to provide semantic hints based on
-        the input description. Uses hed-lsp CLI for tag suggestions.
+        the input description. It uses the feedback LLM to extract keywords,
+        then passes those keywords to hed-lsp CLI for tag suggestions.
         Only runs on the first iteration.
 
         Args:
@@ -176,33 +232,55 @@ class HedAnnotationWorkflow:
 
         logger.info("[WORKFLOW] Entering semantic_preprocess node")
 
-        # Use hed-lsp CLI to suggest tags from the description
-        semantic_hints = []
-        keywords: list[str] = []
+        # Step 1: Extract keywords from the description using LLM
+        keywords = await self._extract_keywords(state["input_description"])
+
+        # Step 2: Use hed-lsp CLI to get tag suggestions for each keyword
+        semantic_hints: list[dict] = []
 
-        if self.hed_lsp_client:
+        if keywords and self.hed_lsp_client:
             try:
-                # Get tag suggestions directly from the description
-                result = self.hed_lsp_client.suggest(state["input_description"])
-                if result.success:
-                    semantic_hints = [
-                        {
-                            "tag": s.tag,
-                            "prefix": "",  # hed-lsp returns full tags
-                            "score": s.score or 0.0,
-                            "source": "hed-lsp",
-                        }
-                        for s in result.suggestions
-                    ]
-                    # Extract keywords from the tags for logging
-                    keywords = [s.tag.split("/")[-1] for s in result.suggestions]
-                    logger.info(
-                        f"[WORKFLOW] hed-lsp suggested {len(semantic_hints)} tags: {keywords[:5]}..."
-                    )
-                else:
-                    logger.warning(f"[WORKFLOW] hed-lsp suggestion failed: {result.error}")
+                # Query hed-lsp for each keyword individually for better results
+                for keyword in keywords:
+                    result = self.hed_lsp_client.suggest(keyword)
+                    if result.success:
+                        for s in result.suggestions:
+                            semantic_hints.append(
+                                {
+                                    "tag": s.tag,
+                                    "keyword": keyword,
+                                    "score": s.score or 0.0,
+                                    "source": "hed-lsp",
+                                }
+                            )
+                    else:
+                        logger.debug(
+                            "[WORKFLOW] hed-lsp suggestion failed for '%s': %s",
+                            keyword,
+                            result.error,
+                        )
+
+                # Deduplicate by tag, keeping highest score
+                seen_tags: dict[str, dict] = {}
+                for hint in semantic_hints:
+                    tag = hint["tag"]
+                    if tag not in seen_tags or hint["score"] > seen_tags[tag]["score"]:
+                        seen_tags[tag] = hint
+                semantic_hints = sorted(seen_tags.values(), key=lambda h: h["score"], reverse=True)
+
+                logger.info(
+                    "[WORKFLOW] hed-lsp suggested %d unique tags from %d keywords",
+                    len(semantic_hints),
+                    len(keywords),
+                )
             except Exception as e:
                 logger.warning("[WORKFLOW] hed-lsp error: %s", e, exc_info=True)
+        elif keywords:
+            # LSP not available; still store keywords for the annotation agent
+            logger.info(
+                "[WORKFLOW] hed-lsp not available; storing %d extracted keywords",
+                len(keywords),
+            )
 
         return {
             "extracted_keywords": keywords,

{hedit-0.7.5a2 → hedit-0.7.5.dev2}/src/api/main.py

@@ -649,6 +649,7 @@ async def annotate(
             schema_version=request.schema_version,
             max_validation_attempts=request.max_validation_attempts,
             run_assessment=request.run_assessment,
+            no_extend=request.no_extend,
             config=config,
         )
         latency_ms = int((time.time() - start_time) * 1000)
@@ -856,6 +857,7 @@ async def annotate_from_image(
             schema_version=request.schema_version,
             max_validation_attempts=request.max_validation_attempts,
             run_assessment=request.run_assessment,
+            no_extend=request.no_extend,
             config=config,
         )
         latency_ms = int((time.time() - start_time) * 1000)
@@ -925,6 +927,64 @@ async def annotate_from_image(
         ) from e
 
 
+async def _collect_stream_telemetry(
+    request: AnnotationRequest | ImageAnnotationRequest,
+    req: Request,
+    current_state: dict,
+    start_time: float,
+    source: str,
+    description: str,
+) -> None:
+    """Collect telemetry for streaming endpoints.
+
+    Shared helper used by both /annotate/stream and /annotate-from-image/stream.
+    Silently returns if telemetry is disabled or collector is not initialized.
+
+    Args:
+        request: The annotation request (text or image)
+        req: FastAPI request for header extraction
+        current_state: Current workflow state dict
+        start_time: Workflow start time (from time.time())
+        source: Telemetry source identifier (e.g., "api-stream", "api-image-stream")
+        description: Input description text (or image description for image endpoints)
+    """
+    if not request.telemetry_enabled or not telemetry_collector:
+        return
+
+    latency_ms = int((time.time() - start_time) * 1000)
+
+    # Get model info from request body, BYOK headers, or server config
+    model_name = (
+        request.model
+        or req.headers.get("x-openrouter-model")
+        or os.getenv("ANNOTATION_MODEL", "openai/gpt-oss-120b")
+    )
+    temperature = request.temperature
+    if temperature is None:
+        temp_header = req.headers.get("x-openrouter-temperature")
+        if temp_header is not None:
+            try:
+                temperature = float(temp_header)
+            except ValueError:
+                temperature = None
+    if temperature is None:
+        temperature = _byok_config.get("temperature", 0.1)
+
+    event = TelemetryEvent.create(
+        description=description,
+        schema_version=request.schema_version,
+        hed_string=current_state.get("current_annotation", ""),
+        iterations=current_state.get("validation_attempts", 0),
+        validation_errors=current_state.get("validation_errors", []),
+        model=model_name,
+        provider=request.provider or req.headers.get("x-openrouter-provider"),
+        temperature=temperature,
+        latency_ms=latency_ms,
+        source=source,
+    )
+    await telemetry_collector.collect(event)
+
+
 @app.post("/annotate/stream")
 async def annotate_stream(
     request: AnnotationRequest,
@@ -1019,6 +1079,7 @@ async def annotate_stream(
         request.schema_version,
         request.max_validation_attempts,
         run_assessment=request.run_assessment,
+        no_extend=request.no_extend,
     )
 
     # Node name to user-friendly stage mapping
@@ -1039,6 +1100,9 @@ async def annotate_stream(
         # SSE padding comment to force Safari to open the stream
         yield ": stream opened\n\n"
 
+        start_time = time.time()
+        current_state = initial_state.copy()
+
         try:
             # Send initial start event
             yield send_event(
@@ -1046,7 +1110,6 @@ async def annotate_stream(
             )
 
             # Track state and progress
-            current_state = initial_state.copy()
             last_stage = None
             validation_attempt = 0
 
@@ -1129,6 +1192,20 @@ async def annotate_stream(
             }
 
             yield send_event("result", result)
+
+            # Collect telemetry after sending result but before done event
+            try:
+                await _collect_stream_telemetry(
+                    request=request,
+                    req=req,
+                    current_state=current_state,
+                    start_time=start_time,
+                    source="api-stream",
+                    description=request.description,
+                )
+            except Exception:
+                logging.debug("Telemetry collection failed for streaming request", exc_info=True)
+
             yield send_event("done", {"message": "Workflow completed"})
 
         except asyncio.CancelledError:
@@ -1142,6 +1219,18 @@ async def annotate_stream(
                     "error_type": "timeout",
                 },
             )
+            # Collect telemetry on error
+            try:
+                await _collect_stream_telemetry(
+                    request=request,
+                    req=req,
+                    current_state=current_state,
+                    start_time=start_time,
+                    source="api-stream",
+                    description=request.description,
+                )
+            except Exception:
+                logging.debug("Telemetry collection failed on timeout", exc_info=True)
             yield send_event("done", {"message": "Workflow ended with error"})
         except RateLimitError:
             logging.exception("Streaming workflow rate limit")
@@ -1152,6 +1241,18 @@ async def annotate_stream(
                     "error_type": "rate_limit",
                 },
             )
+            # Collect telemetry on error
+            try:
+                await _collect_stream_telemetry(
+                    request=request,
+                    req=req,
+                    current_state=current_state,
+                    start_time=start_time,
+                    source="api-stream",
+                    description=request.description,
+                )
+            except Exception:
+                logging.debug("Telemetry collection failed on rate limit", exc_info=True)
             yield send_event("done", {"message": "Workflow ended with error"})
         except Exception:
             logging.exception("Streaming workflow error")
@@ -1162,6 +1263,18 @@ async def annotate_stream(
                     "error_type": "internal",
                 },
             )
+            # Collect telemetry on error
+            try:
+                await _collect_stream_telemetry(
+                    request=request,
+                    req=req,
+                    current_state=current_state,
+                    start_time=start_time,
+                    source="api-stream",
+                    description=request.description,
+                )
+            except Exception:
+                logging.debug("Telemetry collection failed on error", exc_info=True)
             yield send_event("done", {"message": "Workflow ended with error"})
 
     return StreamingResponse(
@@ -1306,6 +1419,10 @@ async def annotate_from_image_stream(
         # SSE padding comment to force Safari to open the stream
         yield ": stream opened\n\n"
 
+        start_time = time.time()
+        current_state: dict = {}
+        image_description = ""
+
         try:
             # Send initial start event
             yield send_event(
@@ -1337,6 +1454,7 @@ async def annotate_from_image_stream(
                 request.schema_version,
                 request.max_validation_attempts,
                 run_assessment=request.run_assessment,
+                no_extend=request.no_extend,
             )
 
             # Track state and progress
@@ -1425,6 +1543,22 @@ async def annotate_from_image_stream(
             }
 
             yield send_event("result", result)
+
+            # Collect telemetry after sending result but before done event
+            try:
+                await _collect_stream_telemetry(
+                    request=request,
+                    req=req,
+                    current_state=current_state,
+                    start_time=start_time,
+                    source="api-image-stream",
+                    description=image_description,
+                )
+            except Exception:
+                logging.debug(
+                    "Telemetry collection failed for image streaming request", exc_info=True
+                )
+
             yield send_event("done", {"message": "Workflow completed"})
 
         except asyncio.CancelledError:
@@ -1438,6 +1572,18 @@ async def annotate_from_image_stream(
                     "error_type": "timeout",
                 },
             )
+            # Collect telemetry on error
+            try:
+                await _collect_stream_telemetry(
+                    request=request,
+                    req=req,
+                    current_state=current_state,
+                    start_time=start_time,
+                    source="api-image-stream",
+                    description=image_description or "image-annotation-failed",
+                )
+            except Exception:
+                logging.debug("Telemetry collection failed on image timeout", exc_info=True)
             yield send_event("done", {"message": "Workflow ended with error"})
         except RateLimitError:
             logging.exception("Streaming image workflow rate limit")
@@ -1448,6 +1594,18 @@ async def annotate_from_image_stream(
                     "error_type": "rate_limit",
                 },
             )
+            # Collect telemetry on error
+            try:
+                await _collect_stream_telemetry(
+                    request=request,
+                    req=req,
+                    current_state=current_state,
+                    start_time=start_time,
+                    source="api-image-stream",
+                    description=image_description or "image-annotation-failed",
+                )
+            except Exception:
+                logging.debug("Telemetry collection failed on image rate limit", exc_info=True)
             yield send_event("done", {"message": "Workflow ended with error"})
         except Exception:
             logging.exception("Streaming image annotation workflow error")
@@ -1458,6 +1616,18 @@ async def annotate_from_image_stream(
                     "error_type": "internal",
                 },
             )
+            # Collect telemetry on error
+            try:
+                await _collect_stream_telemetry(
+                    request=request,
+                    req=req,
+                    current_state=current_state,
+                    start_time=start_time,
+                    source="api-image-stream",
+                    description=image_description or "image-annotation-failed",
+                )
+            except Exception:
+                logging.debug("Telemetry collection failed on image error", exc_info=True)
             yield send_event("done", {"message": "Workflow ended with error"})
 
     return StreamingResponse(

{hedit-0.7.5a2 → hedit-0.7.5.dev2}/src/api/models.py

@@ -55,6 +55,10 @@ class AnnotationRequest(BaseModel):
         le=1.0,
         examples=[0.1, 0.3, 0.7],
     )
+    no_extend: bool = Field(
+        default=False,
+        description="If True, prohibit tag extensions (use only existing HED vocabulary)",
+    )
     telemetry_enabled: bool = Field(
         default=True,
         description="Allow telemetry collection for this request",
@@ -187,6 +191,10 @@ class ImageAnnotationRequest(BaseModel):
         le=1.0,
         examples=[0.1, 0.3, 0.7],
     )
+    no_extend: bool = Field(
+        default=False,
+        description="If True, prohibit tag extensions (use only existing HED vocabulary)",
+    )
     telemetry_enabled: bool = Field(
         default=True,
         description="Allow telemetry collection for this request",

{hedit-0.7.5a2 → hedit-0.7.5.dev2}/src/cli/api_executor.py

@@ -87,6 +87,7 @@ class APIExecutionBackend(ExecutionBackend):
         schema_version: str = "8.4.0",
         max_validation_attempts: int = 5,
         run_assessment: bool = False,
+        no_extend: bool = False,
         **kwargs: Any,
     ) -> dict[str, Any]:
         """Generate HED annotation via API."""
@@ -96,6 +97,7 @@ class APIExecutionBackend(ExecutionBackend):
                 schema_version=schema_version,
                 max_validation_attempts=max_validation_attempts,
                 run_assessment=run_assessment,
+                no_extend=no_extend,
             )
         except APIError as e:
             raise ExecutionError(
@@ -110,6 +112,7 @@ class APIExecutionBackend(ExecutionBackend):
         schema_version: str = "8.4.0",
         max_validation_attempts: int = 5,
         run_assessment: bool = False,
+        no_extend: bool = False,
         **kwargs: Any,
     ) -> Generator[tuple[str, dict[str, Any]], None, None]:
         """Generate HED annotation with streaming progress via API.
@@ -122,6 +125,7 @@ class APIExecutionBackend(ExecutionBackend):
                 schema_version=schema_version,
                 max_validation_attempts=max_validation_attempts,
                 run_assessment=run_assessment,
+                no_extend=no_extend,
             )
         except APIError as e:
             raise ExecutionError(
@@ -137,6 +141,7 @@ class APIExecutionBackend(ExecutionBackend):
         schema_version: str = "8.4.0",
         max_validation_attempts: int = 5,
         run_assessment: bool = False,
+        no_extend: bool = False,
         **kwargs: Any,
     ) -> dict[str, Any]:
         """Generate HED annotation from image via API."""
@@ -147,6 +152,7 @@ class APIExecutionBackend(ExecutionBackend):
                 schema_version=schema_version,
                 max_validation_attempts=max_validation_attempts,
                 run_assessment=run_assessment,
+                no_extend=no_extend,
             )
         except APIError as e:
             raise ExecutionError(
@@ -162,6 +168,7 @@ class APIExecutionBackend(ExecutionBackend):
         schema_version: str = "8.4.0",
         max_validation_attempts: int = 5,
         run_assessment: bool = False,
+        no_extend: bool = False,
         **kwargs: Any,
     ) -> Generator[tuple[str, dict[str, Any]], None, None]:
         """Generate HED annotation from image with streaming progress via API.
@@ -175,6 +182,7 @@ class APIExecutionBackend(ExecutionBackend):
                 schema_version=schema_version,
                 max_validation_attempts=max_validation_attempts,
                 run_assessment=run_assessment,
+                no_extend=no_extend,
             )
         except APIError as e:
             raise ExecutionError(