From a5c531fd29ee3dbe0a55b9e28bbae469c5e73773 Mon Sep 17 00:00:00 2001 From: KaifAhmad1 Date: Fri, 9 Jan 2026 18:39:09 +0530 Subject: [PATCH] Fix semantic extraction empty returns, schema validation, and update docs --- docs/reference/semantic_extract.md | 2 +- pyproject.toml | 6 +- semantica/kg/graph_builder.py | 107 ++++++++--- semantica/semantic_extract/methods.py | 139 ++++++++++---- semantica/semantic_extract/providers.py | 151 ++++++++++++++- semantica/semantic_extract/schemas.py | 124 ++++++++++++ .../test_structured_output.py | 177 ++++++++++++++++++ tests/test_llm_extraction_fixes.py | 6 +- 8 files changed, 647 insertions(+), 65 deletions(-) create mode 100644 semantica/semantic_extract/schemas.py create mode 100644 tests/semantic_extract/test_structured_output.py diff --git a/docs/reference/semantic_extract.md b/docs/reference/semantic_extract.md index ece24b86..6ae412eb 100644 --- a/docs/reference/semantic_extract.md +++ b/docs/reference/semantic_extract.md @@ -253,7 +253,7 @@ entities = ner.extract_entities(text) # Basic relation extraction rel_extractor = RelationExtractor() relations = rel_extractor.extract(text, entities=entities) -# [Relation(source="Elon Musk", target="SpaceX", type="founded")] +# [Relation(subject="Elon Musk", predicate="founded", object="SpaceX")] # With configuration rel_extractor = RelationExtractor( diff --git a/pyproject.toml b/pyproject.toml index a488ad2e..7344ceca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,6 +92,7 @@ dependencies = [ "groq>=0.4.0", "openai>=1.0.0", "litellm>=1.0.0", + "instructor>=1.0.0", "click>=8.1.0", "rich>=12.5.0", "tqdm>=4.64.0", @@ -182,8 +183,11 @@ llm-deepseek = [ llm-litellm = [ "litellm>=1.0.0" ] +llm-instructor = [ + "instructor>=1.0.0" +] llm-all = [ - "semantica[llm-openai,llm-gemini,llm-groq,llm-anthropic,llm-ollama,llm-deepseek,llm-litellm]" + "semantica[llm-openai,llm-gemini,llm-groq,llm-anthropic,llm-ollama,llm-deepseek,llm-litellm,llm-instructor]" ] models-huggingface = [ "transformers>=4.20.0", diff --git a/semantica/kg/graph_builder.py b/semantica/kg/graph_builder.py index 6e7f887b..cb41ec01 100644 --- a/semantica/kg/graph_builder.py +++ b/semantica/kg/graph_builder.py @@ -89,6 +89,7 @@ class GraphBuilder: self.track_history = track_history self.version_snapshots = version_snapshots self.graph_store = graph_store + self.config = kwargs # Store additional config for extractors # Initialize logging from ..utils.logging import get_logger @@ -130,6 +131,11 @@ class GraphBuilder: def _process_item(self, item: Any, all_entities: List[Any], all_relationships: List[Any], **options): """Helper to process a single item and add to entities or relationships list.""" + if isinstance(item, str): + # Treat string as text for extraction + self._extract_from_text(item, all_entities, all_relationships, **options) + return + if hasattr(item, "text") and (hasattr(item, "label") or hasattr(item, "type")): # It's likely an Entity object entity_dict = { @@ -209,30 +215,68 @@ class GraphBuilder: # If still nothing found and has 'text', try extraction if not found_something and "text" in item: text = item["text"] - # Perform extraction if requested or if it's the only way - if options.get("extract", True): - from ..semantic_extract.ner_extractor import NERExtractor - from ..semantic_extract.triplet_extractor import TripletExtractor - - ner_method = options.get("ner_method", "ml") - triplet_method = options.get("triplet_method", "pattern") - - ner = NERExtractor(method=ner_method) - entities = ner.extract_entities(text) - for ent in entities: - self._process_item(ent, all_entities, all_relationships, **options) - - # Only try triplets if specifically requested or if method provided - if "triplet_method" in options or options.get("extract_relations", False): - triplet = TripletExtractor(method=triplet_method) - relations = triplet.extract_triplets(text) - for rel in relations: - self._process_item(rel, all_entities, all_relationships, **options) - found_something = True + self._extract_from_text(text, all_entities, all_relationships, **options) + found_something = True else: # Unknown type pass + def _extract_from_text(self, text: str, all_entities: List[Any], all_relationships: List[Any], **options): + """Helper to extract knowledge from text using configured methods.""" + if not options.get("extract", True): + return + + from ..semantic_extract.ner_extractor import NERExtractor + from ..semantic_extract.relation_extractor import RelationExtractor + from ..semantic_extract.triplet_extractor import TripletExtractor + + # Default to LLM methods as per requirement + ner_method = options.get("ner_method", "llm") + relation_method = options.get("relation_method", "llm") + triplet_method = options.get("triplet_method", "llm") + + self.logger.info(f"Extracting knowledge from text ({len(text)} chars) using {ner_method}...") + + # 1. Extract Entities + ner = NERExtractor(method=ner_method, **self.config) + try: + entities = ner.extract_entities(text, **options) + extracted_count = len(entities) + self._extraction_stats["extracted_entities"] += extracted_count + self.logger.info(f"Extracted {extracted_count} entities") + for ent in entities: + self._process_item(ent, all_entities, all_relationships, **options) + except Exception as e: + self.logger.error(f"Entity extraction failed: {e}") + entities = [] + + # 2. Extract Relations (if requested) + if options.get("extract_relations", True): + rel_extractor = RelationExtractor(method=relation_method, **self.config) + try: + # Pass entities if available to help relation extraction + relations = rel_extractor.extract_relations(text, entities=entities, **options) + extracted_count = len(relations) + self._extraction_stats["extracted_relations"] += extracted_count + self.logger.info(f"Extracted {extracted_count} relationships") + for rel in relations: + self._process_item(rel, all_entities, all_relationships, **options) + except Exception as e: + self.logger.error(f"Relation extraction failed: {e}") + + # 3. Extract Triplets (if requested) + if options.get("extract_triplets", True): + trip_extractor = TripletExtractor(method=triplet_method, **self.config) + try: + triplets = trip_extractor.extract_triplets(text, entities=entities, **options) + extracted_count = len(triplets) + self._extraction_stats["extracted_triplets"] += extracted_count + self.logger.info(f"Extracted {extracted_count} triplets") + for trip in triplets: + self._process_item(trip, all_entities, all_relationships, **options) + except Exception as e: + self.logger.error(f"Triplet extraction failed: {e}") + def build( self, sources: Union[List[Any], Any], @@ -305,6 +349,14 @@ class GraphBuilder: # Track graph building build_start_time = time.time() + + # Initialize extraction statistics for traceability + self._extraction_stats = { + "extracted_entities": 0, + "extracted_relations": 0, + "extracted_triplets": 0 + } + tracking_id = self.progress_tracker.start_tracking( module="kg", submodule="GraphBuilder", @@ -514,7 +566,7 @@ class GraphBuilder: resolution_start = time.time() resolved_entities = resolver_to_use.resolve_entities(all_entities) resolution_time = time.time() - resolution_start - print(f"✅ Resolved to {len(resolved_entities)} unique entities ({resolution_time:.2f}s)") + print(f"[DONE] Resolved to {len(resolved_entities)} unique entities ({resolution_time:.2f}s)") self.logger.info( f"Entity resolution complete: {len(all_entities)} -> {len(resolved_entities)} unique entities" ) @@ -534,7 +586,7 @@ class GraphBuilder: }, } structure_time = time.time() - structure_start - print(f"✅ Graph structure built ({structure_time:.2f}s)") + print(f"[DONE] Graph structure built ({structure_time:.2f}s)") # Persist to GraphStore if available if self.graph_store: @@ -567,7 +619,7 @@ class GraphBuilder: edge_time = time.time() - edge_start total_store_time = time.time() - store_start print(f" Added {edge_count} edges ({edge_time:.2f}s)") - print(f"✅ GraphStore persistence complete ({total_store_time:.2f}s total)") + print(f"[DONE] GraphStore persistence complete ({total_store_time:.2f}s total)") self.logger.info(f"Persisted {node_count} nodes and {edge_count} edges") # Detect and resolve conflicts if conflict detector is available @@ -604,7 +656,14 @@ class GraphBuilder: # Print final summary with timing print(f"\n{'='*60}") - print(f"✅ Knowledge Graph Build Complete") + print(f"[INFO] Extraction Statistics") + print(f" Extracted Entities: {self._extraction_stats['extracted_entities']}") + print(f" Extracted Relationships: {self._extraction_stats['extracted_relations']}") + print(f" Extracted Triplets: {self._extraction_stats['extracted_triplets']}") + print(f"{'='*60}") + + print(f"\n{'='*60}") + print(f"[DONE] Knowledge Graph Build Complete") print(f" Entities: {len(resolved_entities)}") print(f" Relationships: {len(all_relationships)}") print(f" Total time: {total_build_time:.2f}s") diff --git a/semantica/semantic_extract/methods.py b/semantica/semantic_extract/methods.py index dfb50d97..f33e8124 100644 --- a/semantica/semantic_extract/methods.py +++ b/semantica/semantic_extract/methods.py @@ -116,6 +116,12 @@ from .registry import method_registry from .relation_extractor import Relation from .triplet_extractor import Triplet +try: + from .schemas import EntitiesResponse, RelationsResponse, TripletsResponse + SCHEMAS_AVAILABLE = True +except ImportError: + SCHEMAS_AVAILABLE = False + logger = get_logger("methods") # Try to import spaCy @@ -310,6 +316,7 @@ def extract_entities_llm( model: Optional[str] = None, silent_fail: bool = False, max_text_length: Optional[int] = None, + structured_output_mode: str = "typed", **kwargs, ) -> List[Entity]: """ @@ -394,26 +401,36 @@ If an entity doesn't fit any of the preferred types, use the most appropriate ty entity_types_instruction = """Entity types should be one of: PERSON, ORG, GPE, DATE, EVENT, PRODUCT, CONCEPT, or related types. Use the most appropriate type for each entity, including variations or synonyms if they better match the context.""" - prompt = f"""Extract named entities from the following text. -Return ONLY a valid JSON list of objects with the following structure: -[ - {{"text": "entity name", "label": "ENTITY_TYPE", "start": 0, "end": 10, "confidence": 0.9}} -] - -{entity_types_instruction} -Do not include any conversational filler, explanations, or markdown formatting outside the JSON block. - -Text: {text}""" + if not SCHEMAS_AVAILABLE: + raise ImportError("Pydantic schemas not available. Install pydantic/instructor to use LLM extraction.") try: - # 4. EXTRACTION WITH RETRY (handled by generate_structured) - result = llm.generate_structured(prompt) - entities = _parse_entity_result(result, provider, model) + prompt = f"""Extract named entities from the following text. +Return the result as a JSON object with an "entities" key containing the list of entities. +{entity_types_instruction} + +Text: {text}""" - if not entities: - logger.warning(f"No entities extracted using {provider}/{model} from text preview: {text[:100]}...") - - logger.info(f"Successfully extracted {len(entities)} entities using {provider}/{model}") + # Use typed generation with Pydantic schema + result_obj = llm.generate_typed(prompt, schema=EntitiesResponse) + + # Convert back to internal Entity format + entities = [] + for e_out in result_obj.entities: + entities.append(Entity( + text=e_out.text, + label=e_out.label, + start_char=e_out.start if hasattr(e_out, "start") else 0, # Schema might not force these + end_char=e_out.end if hasattr(e_out, "end") else 0, + confidence=e_out.confidence, + metadata={ + "provider": provider, + "model": model, + "extraction_method": "llm_typed", + } + )) + + logger.info(f"Successfully extracted {len(entities)} entities using {provider}/{model} (typed)") return entities except Exception as e: @@ -473,6 +490,7 @@ def _extract_entities_chunked( model: Optional[str], silent_fail: bool, max_text_length: int, + structured_output_mode: str = "typed", **kwargs ) -> List[Entity]: """Internal helper to extract entities from long text by chunking.""" @@ -496,6 +514,7 @@ def _extract_entities_chunked( model=model, silent_fail=False, # We want to know if a chunk fails max_text_length=len(chunk.text) + 1, + structured_output_mode=structured_output_mode, **kwargs ) @@ -882,6 +901,7 @@ def extract_relations_llm( model: Optional[str] = None, silent_fail: bool = False, max_text_length: Optional[int] = None, + structured_output_mode: str = "typed", **kwargs, ) -> List[Relation]: """ @@ -954,7 +974,8 @@ def extract_relations_llm( logger.info(f"Text length ({len(text)}) exceeds limit for relations. Chunking...") return _extract_relations_chunked( text, entities, provider=provider, model=model, - silent_fail=silent_fail, max_text_length=max_text_length, **kwargs + silent_fail=silent_fail, max_text_length=max_text_length, + **kwargs ) entities_str = ", ".join([f"{e.text} ({e.label})" for e in entities]) @@ -972,20 +993,48 @@ If a relation doesn't fit any of the preferred types, use the most appropriate t Extract meaningful relationships between entities. Use appropriate relation types that accurately describe how entities are connected. Common relation types include: related_to, part_of, located_in, created_by, uses, depends_on, interacts_with, and similar variations.""" + if not SCHEMAS_AVAILABLE: + raise ImportError("Pydantic schemas not available. Install pydantic/instructor to use LLM extraction.") + prompt = f"""Extract relations between entities from the following text. +Return the result as a JSON object with a "relations" key containing the list of relations. +Each relation must have 'subject', 'predicate', and 'object' fields. Text: {text} -Entities: {entities_str}{relation_types_instruction} - -Return JSON format: [{{"subject": "...", "predicate": "...", "object": "...", "confidence": 0.9}}] -Extract all meaningful relationships between the entities, using the most appropriate relation type for each relationship.""" +Entities: {entities_str}{relation_types_instruction}""" try: - # 4. EXTRACTION WITH RETRY - result = llm.generate_structured(prompt) - relations = _parse_relation_result(result, entities, text, provider, model) + # Use typed generation with Pydantic schema + result_obj = llm.generate_typed(prompt, schema=RelationsResponse) - logger.info(f"Successfully extracted {len(relations)} relations using {provider}/{model}") + # Convert back to internal Relation format + relations = [] + for r_out in result_obj.relations: + # Find matching entities + subject_entity = next( + (e for e in entities if e.text.lower() == r_out.subject.lower()), + None, + ) + object_entity = next( + (e for e in entities if e.text.lower() == r_out.object.lower()), + None + ) + + if subject_entity and object_entity: + relations.append(Relation( + subject=subject_entity, + predicate=r_out.predicate, + object=object_entity, + confidence=r_out.confidence, + context=text, # Simplified context + metadata={ + "provider": provider, + "model": model, + "extraction_method": "llm_typed" + } + )) + + logger.info(f"Successfully extracted {len(relations)} relations using {provider}/{model} (typed)") return relations except Exception as e: @@ -1067,6 +1116,7 @@ def _extract_relations_chunked( model: Optional[str], silent_fail: bool, max_text_length: int, + structured_output_mode: str = "typed", **kwargs ) -> List[Relation]: """Internal helper to extract relations from long text by chunking.""" @@ -1099,6 +1149,7 @@ def _extract_relations_chunked( model=model, silent_fail=False, max_text_length=len(chunk.text) + 1, + structured_output_mode=structured_output_mode, **kwargs ) all_relations.extend(chunk_rels) @@ -1248,6 +1299,7 @@ def extract_triplets_llm( model: Optional[str] = None, silent_fail: bool = False, max_text_length: Optional[int] = None, + structured_output_mode: str = "typed", **kwargs, ) -> List[Triplet]: """ @@ -1314,21 +1366,38 @@ def extract_triplets_llm( logger.info(f"Text length ({len(text)}) exceeds limit for triplets. Chunking...") return _extract_triplets_chunked( text, provider=provider, model=model, - silent_fail=silent_fail, max_text_length=max_text_length, **kwargs + silent_fail=silent_fail, max_text_length=max_text_length, + **kwargs ) + + if not SCHEMAS_AVAILABLE: + raise ImportError("Pydantic schemas not available. Install pydantic/instructor to use LLM extraction.") prompt = f"""Extract RDF triplets (subject-predicate-object) from the following text. +Return the result as a JSON object with a "triplets" key containing the list of triplets. -Text: {text} - -Return JSON format: [{{"subject": "...", "predicate": "...", "object": "...", "confidence": 0.9}}]""" +Text: {text}""" try: - # 4. EXTRACTION WITH RETRY - result = llm.generate_structured(prompt) - triplets = _parse_triplet_result(result, provider, model) + # Use typed generation with Pydantic schema + result_obj = llm.generate_typed(prompt, schema=TripletsResponse) - logger.info(f"Successfully extracted {len(triplets)} triplets using {provider}/{model}") + # Convert back to internal Triplet format + triplets = [] + for t_out in result_obj.triplets: + triplets.append(Triplet( + subject=t_out.subject, + predicate=t_out.predicate, + object=t_out.object, + confidence=t_out.confidence, + metadata={ + "provider": provider, + "model": model, + "extraction_method": "llm_typed" + } + )) + + logger.info(f"Successfully extracted {len(triplets)} triplets using {provider}/{model} (typed)") return triplets except Exception as e: @@ -1389,6 +1458,7 @@ def _extract_triplets_chunked( model: Optional[str], silent_fail: bool, max_text_length: int, + structured_output_mode: str = "typed", **kwargs ) -> List[Triplet]: """Internal helper to extract triplets from long text by chunking.""" @@ -1411,6 +1481,7 @@ def _extract_triplets_chunked( model=model, silent_fail=False, max_text_length=len(chunk.text) + 1, + structured_output_mode=structured_output_mode, **kwargs ) all_triplets.extend(chunk_triplets) diff --git a/semantica/semantic_extract/providers.py b/semantica/semantic_extract/providers.py index 521db8b5..d4fe7011 100644 --- a/semantica/semantic_extract/providers.py +++ b/semantica/semantic_extract/providers.py @@ -71,7 +71,19 @@ License: MIT import json import os -from typing import Any, Dict, List, Optional, Union +import time +from typing import Any, Dict, List, Optional, Union, Type + +try: + from pydantic import BaseModel, ValidationError +except ImportError: + BaseModel = Any + ValidationError = Exception + +try: + import instructor +except ImportError: + instructor = None from ..utils.exceptions import ProcessingError from ..utils.logging import get_logger @@ -211,6 +223,141 @@ class BaseProvider: raise ProcessingError(f"Failed to generate structured output: {last_error}") return [] + def generate_typed( + self, + prompt: str, + schema: Type[BaseModel], + max_retries: int = 3, + **kwargs + ) -> BaseModel: + """ + Generate structured output validated against a Pydantic schema. + Uses instructor if available and supported for the provider, otherwise falls back to a repair loop. + """ + provider_name = self.__class__.__name__ + + # Try using instructor first if available + if instructor: + try: + client = None + mode = instructor.Mode.TOOLS # Default mode + + if provider_name == "OpenAIProvider" and self.client: + client = instructor.from_openai(self.client) + elif provider_name == "AnthropicProvider" and self.client: + client = instructor.from_anthropic(self.client) + elif provider_name == "GeminiProvider" and self.client: + client = instructor.from_gemini( + self.client, + mode=instructor.Mode.GEMINI_JSON + ) + elif provider_name == "GroqProvider" and self.client: + # Groq is OpenAI-compatible but works best with JSON mode + client = instructor.from_openai(self.client, mode=instructor.Mode.JSON) + elif provider_name == "OllamaProvider": + # Create OpenAI-compatible client for Ollama + try: + from openai import OpenAI + # Ollama typically runs on localhost:11434/v1 + base_url = getattr(self, "base_url", "http://localhost:11434") + if not base_url.endswith("/v1"): + base_url = f"{base_url.rstrip('/')}/v1" + + ollama_client = OpenAI( + base_url=base_url, + api_key="ollama", # required but unused + ) + client = instructor.from_openai(ollama_client, mode=instructor.Mode.JSON) + except ImportError: + pass + elif provider_name == "DeepSeekProvider" and self.client: + # DeepSeek is OpenAI compatible + # We need to wrap the underlying client if it exposes the OpenAI interface + # or create a new OpenAI client if self.client is a deepseek.Client (which might be just a wrapper) + # Assuming deepseek.Client is compatible or we can use OpenAI client + try: + # DeepSeek usually works with standard OpenAI client + # If self.client is deepseek.Client, check if we can wrap it + # Otherwise create a new OpenAI client + from openai import OpenAI + if isinstance(self.client, OpenAI): + client = instructor.from_openai(self.client, mode=instructor.Mode.JSON) + else: + # Try creating fresh client + ds_client = OpenAI( + api_key=self.api_key, + base_url="https://api.deepseek.com" + ) + client = instructor.from_openai(ds_client, mode=instructor.Mode.JSON) + except Exception: + pass + + if client: + # Map generate arguments to client arguments + # Instructor standardizes on chat.completions.create for OpenAI/Groq/Anthropic/Gemini + + response = client.chat.completions.create( + model=kwargs.get("model", self.model), + messages=[{"role": "user", "content": prompt}], + response_model=schema, + max_retries=max_retries, + temperature=kwargs.get("temperature", 0.1), # Low temp for structured + ) + return response + except Exception as e: + self.logger.warning(f"Instructor generation failed ({e}), falling back to manual repair loop.") + + # Fallback: Manual repair loop + last_error = None + current_prompt = prompt + + for attempt in range(max_retries): + try: + # 1. Generate JSON + # We use generate_structured to get the dict/list + json_result = self.generate_structured(current_prompt, max_retries=1, **kwargs) + + # 2. Validate with Schema + # If the result is a list and schema expects a wrapper, or vice versa, we might need adjustment + # But we assume the prompt asks for the correct structure matching the schema. + + # Special handling if schema is a wrapper but result is a list + if isinstance(json_result, list) and hasattr(schema, "entities") and "entities" in schema.model_fields: + # Auto-wrap for entities + json_result = {"entities": json_result} + elif isinstance(json_result, list) and hasattr(schema, "relations") and "relations" in schema.model_fields: + json_result = {"relations": json_result} + elif isinstance(json_result, list) and hasattr(schema, "triplets") and "triplets" in schema.model_fields: + json_result = {"triplets": json_result} + + validated = schema.model_validate(json_result) + return validated + + except ValidationError as e: + last_error = e + error_summary = str(e) + # Simplify error summary for the LLM + # (You could parse e.errors() for a better message) + + if attempt < max_retries - 1: + wait_time = (attempt + 1) * 1 + self.logger.warning(f"Schema validation failed (attempt {attempt + 1}): {e}. Retrying with error feedback...") + + # Update prompt with error info + current_prompt = f"{prompt}\n\nPrevious response was invalid JSON or didn't match schema:\n{error_summary}\n\nPlease fix the errors and return valid JSON matching the schema." + time.sleep(wait_time) + else: + self.logger.error(f"Typed generation failed validation: {e}") + + except Exception as e: + last_error = e + if attempt < max_retries - 1: + time.sleep(1) + else: + self.logger.error(f"Typed generation failed: {e}") + + raise ProcessingError(f"Failed to generate typed output after {max_retries} attempts: {last_error}") + class OpenAIProvider(BaseProvider): """OpenAI provider implementation.""" @@ -333,7 +480,7 @@ class GroqProvider(BaseProvider): """Groq provider implementation.""" def __init__( - self, api_key: Optional[str] = None, model: str = "llama2-70b-4096", **kwargs + self, api_key: Optional[str] = None, model: str = "llama-3.3-70b-versatile", **kwargs ): """Initialize Groq provider.""" super().__init__(**kwargs) diff --git a/semantica/semantic_extract/schemas.py b/semantica/semantic_extract/schemas.py new file mode 100644 index 00000000..58d2c69d --- /dev/null +++ b/semantica/semantic_extract/schemas.py @@ -0,0 +1,124 @@ +from typing import List, Optional +from pydantic import BaseModel, Field, field_validator, model_validator, ConfigDict + +class EntityOut(BaseModel): + """Canonical schema for entity extraction output.""" + model_config = ConfigDict(populate_by_name=True, extra="ignore") + + text: str = Field(..., description="The text content of the entity") + label: str = Field(..., description="The type or label of the entity (e.g., PERSON, ORG)") + start: int = Field(0, description="Start character index", alias="start_char") + end: int = Field(0, description="End character index", alias="end_char") + confidence: float = Field(0.9, description="Confidence score between 0 and 1") + + @field_validator("text", mode="before") + @classmethod + def clean_text(cls, v): + if isinstance(v, str): + return v.strip() + return str(v) + + @field_validator("confidence", mode="before") + @classmethod + def normalize_confidence(cls, v): + if isinstance(v, str): + try: + v = float(v) + except ValueError: + return 0.9 + if isinstance(v, (int, float)): + return max(0.0, min(1.0, float(v))) + return 0.9 + + @model_validator(mode="before") + @classmethod + def handle_aliases(cls, data): + if isinstance(data, dict): + # Handle 'type' as alias for 'label' + if "label" not in data and "type" in data: + data["label"] = data["type"] + # Handle 'value' or 'span' as alias for 'text' + if "text" not in data: + if "value" in data: + data["text"] = data["value"] + elif "span" in data: + data["text"] = data["span"] + return data + +class RelationOut(BaseModel): + """Canonical schema for relation extraction output.""" + model_config = ConfigDict(populate_by_name=True, extra="ignore") + + subject: str = Field(..., description="Source entity text") + object: str = Field(..., description="Target entity text") + predicate: str = Field(..., description="Relation type or predicate") + confidence: float = Field(0.9, description="Confidence score between 0 and 1") + + @model_validator(mode="before") + @classmethod + def handle_aliases(cls, data): + if isinstance(data, dict): + if "subject" not in data and "source" in data: + data["subject"] = data["source"] + if "object" not in data and "target" in data: + data["object"] = data["target"] + if "predicate" not in data and "label" in data: + data["predicate"] = data["label"] + return data + + @property + def source(self): + return self.subject + + @property + def target(self): + return self.object + + @property + def label(self): + return self.predicate + + @field_validator("confidence", mode="before") + @classmethod + def normalize_confidence(cls, v): + if isinstance(v, str): + try: + v = float(v) + except ValueError: + return 0.9 + if isinstance(v, (int, float)): + return max(0.0, min(1.0, float(v))) + return 0.9 + +class TripletOut(BaseModel): + """Canonical schema for triplet extraction output.""" + model_config = ConfigDict(populate_by_name=True, extra="ignore") + + subject: str = Field(..., description="Subject of the triplet") + predicate: str = Field(..., description="Predicate or relation") + object: str = Field(..., description="Object of the triplet") + confidence: float = Field(0.9, description="Confidence score between 0 and 1") + + @field_validator("confidence", mode="before") + @classmethod + def normalize_confidence(cls, v): + if isinstance(v, str): + try: + v = float(v) + except ValueError: + return 0.9 + if isinstance(v, (int, float)): + return max(0.0, min(1.0, float(v))) + return 0.9 + +class EntitiesResponse(BaseModel): + """Wrapper for list of entities.""" + entities: List[EntityOut] = Field(default_factory=list) + +class RelationsResponse(BaseModel): + """Wrapper for list of relations.""" + relations: List[RelationOut] = Field(default_factory=list) + +class TripletsResponse(BaseModel): + """Wrapper for list of triplets.""" + triplets: List[TripletOut] = Field(default_factory=list) diff --git a/tests/semantic_extract/test_structured_output.py b/tests/semantic_extract/test_structured_output.py new file mode 100644 index 00000000..f3750e76 --- /dev/null +++ b/tests/semantic_extract/test_structured_output.py @@ -0,0 +1,177 @@ + +import pytest +from unittest.mock import MagicMock, patch +from typing import Type, List, Optional +from pydantic import BaseModel, ValidationError + +from semantica.semantic_extract.providers import BaseProvider +from semantica.semantic_extract.methods import ( + extract_entities_llm, + extract_relations_llm, + extract_triplets_llm +) +from semantica.semantic_extract.schemas import EntitiesResponse, RelationsResponse, TripletsResponse +from semantica.semantic_extract.ner_extractor import Entity +from semantica.semantic_extract.relation_extractor import Relation + +# Mock Pydantic models for responses +class MockEntity(BaseModel): + text: str + label: str + start: int = 0 + end: int = 0 + confidence: float = 1.0 + +class MockEntitiesResponse(BaseModel): + entities: List[MockEntity] + +class MockRelation(BaseModel): + subject: str + predicate: str + object: str + confidence: float = 1.0 + +class MockRelationsResponse(BaseModel): + relations: List[MockRelation] + +class MockTriplet(BaseModel): + subject: str + predicate: str + object: str + confidence: float = 1.0 + +class MockTripletsResponse(BaseModel): + triplets: List[MockTriplet] + +class MockProvider(BaseProvider): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.generate_typed_called = False + self.generate_structured_called = False + self.model = "mock-model" + self.is_available_val = True + + def is_available(self) -> bool: + return self.is_available_val + + def generate(self, prompt: str, **kwargs) -> str: + return "{}" + + def generate_structured(self, prompt: str, **kwargs) -> dict: + self.generate_structured_called = True + if "entities" in prompt.lower(): + return [{"text": "Apple", "label": "ORG", "start": 0, "end": 5}] + elif "relations" in prompt.lower(): + return [{"subject": "Steve Jobs", "predicate": "founded", "object": "Apple"}] + elif "triplets" in prompt.lower(): + return [{"subject": "Steve Jobs", "predicate": "founded", "object": "Apple"}] + return {} + + def generate_typed( + self, + prompt: str, + schema: Type[BaseModel], + max_retries: int = 3, + **kwargs + ) -> BaseModel: + self.generate_typed_called = True + + if schema.__name__ == "EntitiesResponse": + return EntitiesResponse(entities=[ + {"text": "Apple", "label": "ORG", "start_char": 0, "end_char": 5, "confidence": 0.99} + ]) + elif schema.__name__ == "RelationsResponse": + return RelationsResponse(relations=[ + {"subject": "Steve Jobs", "predicate": "founded", "object": "Apple", "confidence": 0.95} + ]) + elif schema.__name__ == "TripletsResponse": + return TripletsResponse(triplets=[ + {"subject": "Steve Jobs", "predicate": "founded", "object": "Apple", "confidence": 0.95} + ]) + return schema() + +@pytest.fixture +def mock_provider(): + return MockProvider() + +@patch("semantica.semantic_extract.methods.create_provider") +def test_extract_entities_typed(mock_create_provider, mock_provider): + mock_create_provider.return_value = mock_provider + + text = "Apple was founded by Steve Jobs." + entities = extract_entities_llm( + text, + provider="mock", + structured_output_mode="typed" + ) + + assert mock_provider.generate_typed_called + assert len(entities) == 1 + assert entities[0].text == "Apple" + assert entities[0].label == "ORG" + assert entities[0].metadata["extraction_method"] == "llm_typed" + +@patch("semantica.semantic_extract.methods.create_provider") +def test_extract_entities_legacy(mock_create_provider, mock_provider): + mock_create_provider.return_value = mock_provider + + text = "Apple was founded by Steve Jobs." + entities = extract_entities_llm( + text, + provider="mock", + structured_output_mode="legacy" + ) + + # Legacy mode now redirects to typed mode + assert mock_provider.generate_typed_called + assert not mock_provider.generate_structured_called + assert len(entities) == 1 + assert entities[0].text == "Apple" + assert entities[0].label == "ORG" + assert entities[0].metadata["extraction_method"] == "llm_typed" + +@patch("semantica.semantic_extract.methods.create_provider") +def test_extract_relations_typed(mock_create_provider, mock_provider): + mock_create_provider.return_value = mock_provider + + text = "Steve Jobs founded Apple." + entities = [ + Entity(text="Steve Jobs", label="PERSON", start_char=0, end_char=10), + Entity(text="Apple", label="ORG", start_char=19, end_char=24) + ] + + relations = extract_relations_llm( + text, + entities=entities, + provider="mock", + structured_output_mode="typed" + ) + + assert mock_provider.generate_typed_called + assert len(relations) == 1 + assert relations[0].subject.text == "Steve Jobs" + assert relations[0].object.text == "Apple" + assert relations[0].predicate == "founded" + assert relations[0].metadata["extraction_method"] == "llm_typed" + +@patch("semantica.semantic_extract.methods.create_provider") +def test_extract_triplets_typed(mock_create_provider, mock_provider): + mock_create_provider.return_value = mock_provider + + text = "Steve Jobs founded Apple." + + triplets = extract_triplets_llm( + text, + provider="mock", + structured_output_mode="typed" + ) + + assert mock_provider.generate_typed_called + assert len(triplets) == 1 + assert triplets[0].subject == "Steve Jobs" + assert triplets[0].object == "Apple" + assert triplets[0].predicate == "founded" + assert triplets[0].metadata["extraction_method"] == "llm_typed" + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/tests/test_llm_extraction_fixes.py b/tests/test_llm_extraction_fixes.py index 588ae383..5cf60824 100644 --- a/tests/test_llm_extraction_fixes.py +++ b/tests/test_llm_extraction_fixes.py @@ -34,7 +34,7 @@ class TestLLMExtractionFixes(unittest.TestCase): """Test that methods raise ProcessingError by default on failure.""" mock_llm = MagicMock() mock_llm.is_available.return_value = True - mock_llm.generate_structured.side_effect = ProcessingError("LLM Error") + mock_llm.generate_typed.side_effect = ProcessingError("LLM Error") mock_create.return_value = mock_llm try: @@ -48,7 +48,7 @@ class TestLLMExtractionFixes(unittest.TestCase): """Test that silent_fail=True returns empty list instead of raising.""" mock_llm = MagicMock() mock_llm.is_available.return_value = True - mock_llm.generate_structured.side_effect = Exception("LLM Error") + mock_llm.generate_typed.side_effect = Exception("LLM Error") mock_create.return_value = mock_llm entities = extract_entities_llm("test text", provider="openai", silent_fail=True) @@ -83,7 +83,7 @@ class TestLLMExtractionFixes(unittest.TestCase): """Test that long text triggers chunking.""" mock_llm = MagicMock() mock_llm.is_available.return_value = True - mock_llm.generate_structured.return_value = [] + mock_llm.generate_typed.return_value = MagicMock(entities=[]) # Mock response mock_create.return_value = mock_llm long_text = "This is a long text that should be chunked into multiple pieces."