Jonathan/reasoning mode

langstruct/api.py

@@ -8,6 +8,7 @@
 import dspy
 
 from .core import modules as core_modules
+from .core.modules import ReasoningMode
 from .core.chunking import ChunkingConfig
 from .core.export_utils import ExportUtilities
 from .core.modules import QueryParser
@@ -212,6 +213,8 @@ def extract(
         text: str,
         confidence_threshold: float = 0.0,
         validate: bool = True,
+        run_validation: bool = True,
+        reasoning: Union[ReasoningMode, str] = ReasoningMode.PREDICT,
         debug: bool = False,
         return_sources: Optional[bool] = None,
         refine: Union[bool, Refine, Dict[str, Any], None] = None,
@@ -224,6 +227,8 @@ def extract(
         text: List[str],
         confidence_threshold: float = 0.0,
         validate: bool = True,
+        run_validation: bool = True,
+        reasoning: Union[ReasoningMode, str] = ReasoningMode.PREDICT,
         debug: bool = False,
         return_sources: Optional[bool] = None,
         max_workers: Optional[int] = None,
@@ -239,6 +244,8 @@ def extract(
         text: Union[str, List[str]],
         confidence_threshold: float = 0.0,
         validate: bool = True,
+        run_validation: bool = True,
+        reasoning: Union[ReasoningMode, str] = ReasoningMode.PREDICT,
         debug: bool = False,
         return_sources: Optional[bool] = None,
         max_workers: Optional[int] = None,
@@ -256,7 +263,10 @@ def extract(
         Args:
             text: Input text or list of texts to extract from
             confidence_threshold: Minimum confidence score to accept results
-            validate: Whether to run quality validation and show suggestions
+            validate: Whether to run post-extraction quality validation and show suggestions
+            run_validation: Whether to run the LLM validation step inside the extraction
+                pipeline (i.e., `EntityExtractor.validate`). Set False to skip that call.
+            reasoning: Reasoning strategy used by the extraction pipeline ("predict" or "cot").
             debug: Whether to show detailed validation warnings and suggestions (default: False)
             return_sources: Override source grounding for this call
             max_workers: Maximum parallel workers for batch processing (list input only)
@@ -303,6 +313,8 @@ def extract(
                 texts=text,
                 confidence_threshold=confidence_threshold,
                 validate=validate,
+                run_validation=run_validation,
+                reasoning=reasoning,
                 debug=debug,
                 return_sources=return_sources,
                 max_workers=max_workers,
@@ -314,14 +326,23 @@ def extract(
 
         # Handle single text input
         return self._extract_single(
-            text, confidence_threshold, validate, debug, return_sources, refine
+            text,
+            confidence_threshold,
+            validate,
+            run_validation,
+            reasoning,
+            debug,
+            return_sources,
+            refine,
         )
 
     def _extract_parallel(
         self,
         texts: List[str],
         confidence_threshold: float = 0.0,
         validate: bool = True,
+        run_validation: bool = True,
+        reasoning: Union[ReasoningMode, str] = ReasoningMode.PREDICT,
         debug: bool = False,
         return_sources: Optional[bool] = None,
         max_workers: Optional[int] = None,
@@ -359,6 +380,8 @@ def process_fn(text: str) -> ExtractionResult:
                 text=text,
                 confidence_threshold=confidence_threshold,
                 validate=validate,
+                run_validation=run_validation,
+                reasoning=reasoning,
                 debug=debug,
                 return_sources=return_sources,
                 refine=refine,
@@ -389,12 +412,15 @@ def extract_batch(
         texts: List[str],
         confidence_threshold: float = 0.0,
         validate: bool = True,
+        run_validation: bool = True,
+        reasoning: Union[ReasoningMode, str] = ReasoningMode.PREDICT,
         debug: bool = False,
         return_sources: Optional[bool] = None,
         max_workers: int = 10,
         show_progress: bool = True,
         rate_limit: Optional[int] = None,
         return_failures: bool = False,
+        refine: Union[bool, Refine, Dict[str, Any], None] = None,
     ) -> Union[List[ExtractionResult], ProcessingResult]:
         """Batch extract with explicit parallel processing control.
 
@@ -405,12 +431,16 @@ def extract_batch(
             texts: List of texts to extract from
             confidence_threshold: Minimum confidence score to accept
             validate: Whether to run validation
+            run_validation: Whether to run the LLM validation step inside the extraction
+                pipeline (i.e., `EntityExtractor.validate`). Set False to skip that call.
+            reasoning: Reasoning strategy used by the extraction pipeline ("predict" or "cot").
             debug: Whether to show detailed validation warnings and suggestions (default: False)
             return_sources: Override source grounding
             max_workers: Number of parallel workers (default: 10)
             show_progress: Show progress bar (default: True)
             rate_limit: API calls per minute limit
             return_failures: If True, returns ProcessingResult with successes/failures
+            refine: Override refinement config for this call
 
         Returns:
             List[ExtractionResult] if return_failures=False (raises on any failure)
@@ -437,6 +467,8 @@ def process_fn(text: str) -> ExtractionResult:
                 text=text,
                 confidence_threshold=confidence_threshold,
                 validate=validate,
+                run_validation=run_validation,
+                reasoning=reasoning,
                 debug=debug,
                 return_sources=return_sources,
                 refine=refine,
@@ -462,6 +494,8 @@ def _extract_single(
         text: str,
         confidence_threshold: float = 0.0,
         validate: bool = True,
+        run_validation: bool = True,
+        reasoning: Union[ReasoningMode, str] = ReasoningMode.PREDICT,
         debug: bool = False,
         return_sources: Optional[bool] = None,
         refine: Union[bool, Refine, Dict[str, Any], None] = None,
@@ -484,7 +518,14 @@ def _extract_single(
                 overridden = True
 
             # Run extraction pipeline (call bound __call__ so tests can patch it)
-            result = self.pipeline.__call__(text)
+            # NOTE: some tests monkeypatch `ExtractionPipeline` with mocks that do
+            # not accept `run_validation`/`reasoning`, so fall back gracefully.
+            try:
+                result = self.pipeline.__call__(
+                    text, run_validation=run_validation, reasoning=reasoning
+                )
+            except TypeError:
+                result = self.pipeline.__call__(text)
         finally:
             # Restore previous setting if we overrode it
             if (

langstruct/core/modules.py

@@ -1,7 +1,9 @@
 """Core DSPy extraction modules implementing the extraction pipeline."""
 
 import json
-from typing import Any, Dict, List, Optional, Type
+import logging
+from enum import Enum
+from typing import Any, Dict, List, Optional, Type, Union
 
 import dspy
 from pydantic import BaseModel, ValidationError
@@ -18,35 +20,94 @@
     ValidateExtraction,
 )
 
+logger = logging.getLogger(__name__)
+
+
+class ReasoningMode(str, Enum):
+    """Reasoning strategy for DSPy predictors."""
+
+    COT = "cot"
+    PREDICT = "predict"
+
+
+def _coerce_reasoning_mode(mode: Union["ReasoningMode", str, None]) -> "ReasoningMode":
+    """Coerce a user-provided value into a ReasoningMode."""
+    if mode is None:
+        return ReasoningMode.PREDICT
+    if isinstance(mode, ReasoningMode):
+        return mode
+    try:
+        return ReasoningMode(str(mode).lower())
+    except Exception:
+        return ReasoningMode.PREDICT
+
 
 class EntityExtractor(dspy.Module):
     """Core entity extraction module using DSPy Chain of Thought."""
 
-    def __init__(self, schema: Type[BaseModel], use_sources: bool = True):
+    def __init__(
+        self,
+        schema: Type[BaseModel],
+        use_sources: bool = True,
+        reasoning: Union[ReasoningMode, str] = ReasoningMode.PREDICT,
+    ):
         super().__init__()
         self.schema = schema
         self.use_sources = use_sources
+        self.reasoning = _coerce_reasoning_mode(reasoning)
 
-        # Initialize DSPy modules
+        # Initialize DSPy modules (keep both so we can switch per-call)
         if use_sources:
-            self.extract = dspy.ChainOfThought(ExtractWithSources)
+            self._extract_predict = dspy.Predict(ExtractWithSources)
+            self._extract_cot = dspy.ChainOfThought(ExtractWithSources)
         else:
-            self.extract = dspy.ChainOfThought(ExtractEntities)
+            self._extract_predict = dspy.Predict(ExtractEntities)
+            self._extract_cot = dspy.ChainOfThought(ExtractEntities)
 
-        self.validate = dspy.ChainOfThought(ValidateExtraction)
+        self._validate_predict = dspy.Predict(ValidateExtraction)
+        self._validate_cot = dspy.ChainOfThought(ValidateExtraction)
         self.grounder = SourceGrounder()
 
-    def forward(self, text: str, chunk_offset: int = 0) -> ExtractionResult:
-        """Extract structured entities from text with validation."""
+    def forward(
+        self,
+        text: str,
+        chunk_offset: int = 0,
+        run_validation: bool = True,
+        reasoning: Union[ReasoningMode, str, None] = None,
+    ) -> ExtractionResult:
+        """Extract structured entities from text.
+
+        Args:
+            text: Input text to extract from
+            chunk_offset: Offset to add to source spans (when processing chunks)
+            run_validation: If False, skip the LLM validation step (`self.validate`)
+            reasoning: Override reasoning strategy for this call ("predict" or "cot")
+        """
+        effective_reasoning = (
+            _coerce_reasoning_mode(reasoning)
+            if reasoning is not None
+            else self.reasoning
+        )
+        logger.info(
+            "EntityExtractor.forward reasoning=%s run_validation=%s",
+            effective_reasoning.value,
+            bool(run_validation),
+        )
+
         schema_json = json.dumps(get_json_schema(self.schema), indent=2)
 
         # Perform extraction
+        extractor = (
+            self._extract_cot
+            if effective_reasoning == ReasoningMode.COT
+            else self._extract_predict
+        )
         if self.use_sources:
-            result = self.extract(text=text, schema_spec=schema_json)
+            result = extractor(text=text, schema_spec=schema_json)
             entities_json = result.entities
             sources_json = getattr(result, "sources", "{}")
         else:
-            result = self.extract(text=text, schema_spec=schema_json)
+            result = extractor(text=text, schema_spec=schema_json)
             entities_json = result.entities
             sources_json = "{}"
 
@@ -59,10 +120,22 @@ def forward(self, text: str, chunk_offset: int = 0) -> ExtractionResult:
             entities_dict = {}
             sources_dict = {}
 
-        # Validate extraction using DSPy
-        validation = self.validate(
-            text=text, entities=entities_json, schema_spec=schema_json
-        )
+        # Validate extraction using DSPy (optional)
+        if run_validation:
+            validator = (
+                self._validate_cot
+                if effective_reasoning == ReasoningMode.COT
+                else self._validate_predict
+            )
+            validation = validator(
+                text=text, entities=entities_json, schema_spec=schema_json
+            )
+            is_valid = bool(getattr(validation, "is_valid", True))
+            validation_feedback = getattr(validation, "feedback", "")
+        else:
+            validation = None
+            is_valid = True
+            validation_feedback = "Validation skipped"
 
         # Ground entities to source locations
         if not sources_dict:  # If LLM didn't provide sources, compute them
@@ -80,15 +153,16 @@ def forward(self, text: str, chunk_offset: int = 0) -> ExtractionResult:
             )
 
         # Calculate overall confidence
-        confidence = self._calculate_confidence(validation.is_valid, entities_dict)
+        confidence = self._calculate_confidence(is_valid, entities_dict)
 
         return ExtractionResult(
             entities=entities_dict,
             sources=grounded_sources,
             confidence=confidence,
             metadata={
-                "validation_feedback": validation.feedback,
-                "is_valid": validation.is_valid,
+                "validation_feedback": validation_feedback,
+                "is_valid": is_valid,
+                "validation_skipped": not run_validation,
                 "chunk_offset": chunk_offset,
             },
         )
@@ -407,22 +481,40 @@ def __init__(
         schema: Type[BaseModel],
         chunking_config: Optional[ChunkingConfig] = None,
         use_sources: bool = True,
+        reasoning: Union[ReasoningMode, str] = ReasoningMode.PREDICT,
     ):
         super().__init__()
         self.schema = schema
         self.chunker = TextChunkerModule(chunking_config)
-        self.extractor = EntityExtractor(schema, use_sources)
+        self.extractor = EntityExtractor(schema, use_sources, reasoning=reasoning)
         self.aggregator = ResultAggregator(schema)
 
-    def forward(self, text: str) -> ExtractionResult:
-        """Run the complete extraction pipeline on input text."""
+    def forward(
+        self,
+        text: str,
+        run_validation: bool = True,
+        reasoning: Union[ReasoningMode, str, None] = None,
+    ) -> ExtractionResult:
+        """Run the complete extraction pipeline on input text.
+
+        Args:
+            text: Input text to extract from
+            run_validation: If False, skip the LLM validation step inside
+                `EntityExtractor` (i.e., it will not call `self.validate`).
+            reasoning: Override reasoning strategy for this call ("predict" or "cot")
+        """
         # Step 1: Chunk the text
         chunks = self.chunker(text)
 
         # Step 2: Extract from each chunk
         chunk_results = []
         for chunk in chunks:
-            extraction = self.extractor(chunk.text, chunk.start_offset)
+            extraction = self.extractor(
+                chunk.text,
+                chunk.start_offset,
+                run_validation=run_validation,
+                reasoning=reasoning,
+            )
             chunk_result = ChunkResult(
                 chunk_id=chunk.id,
                 chunk_text=chunk.text,