aiecs 1.0.1__py3-none-any.whl → 1.7.17__py3-none-any.whl

aiecs/tools/task_tools/research_tool.py

@@ -1,46 +1,27 @@
 import logging
-from typing import Dict, Any, List, Optional, Tuple
+from typing import Dict, Any, List, Optional, Set
 import spacy
 from spacy.language import Language
-from pydantic import BaseModel, ValidationError, ConfigDict
-from pydantic_settings import BaseSettings
+from pydantic import BaseModel, Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
 from collections import Counter
-from scipy.stats import pearsonr
+from scipy.stats import pearsonr  # type: ignore[import-untyped]
 import os
 
 from aiecs.tools.base_tool import BaseTool
 from aiecs.tools import register_tool
 
-# Configuration for ResearchTool
-class ResearchSettings(BaseSettings):
-    """
-    Configuration for ResearchTool.
-
-    Attributes:
-        max_workers (int): Maximum number of thread pool workers.
-        spacy_model (str): Default spaCy model to use.
-        max_text_length (int): Maximum text length for inputs.
-        allowed_spacy_models (List[str]): Allowed spaCy models.
-        env_prefix (str): Environment variable prefix.
-    """
-    max_workers: int = min(32, (os.cpu_count() or 4) * 2)
-    spacy_model: str = "en_core_web_sm"
-    max_text_length: int = 10_000
-    allowed_spacy_models: List[str] = ["en_core_web_sm", "zh_core_web_sm"]
-    env_prefix: str = 'RESEARCH_TOOL_'
-
-    model_config = ConfigDict(env_prefix='RESEARCH_TOOL_')
 
 # Exceptions
 class ResearchToolError(Exception):
     """Base exception for ResearchTool errors."""
-    pass
+
 
 class FileOperationError(ResearchToolError):
     """Raised when file operations fail."""
-    pass
 
-@register_tool('research')
+
+@register_tool("research")
 class ResearchTool(BaseTool):
     """
     Tool for causal inference using Mill's methods, advanced induction, deduction, and text summarization.
@@ -57,34 +38,112 @@ class ResearchTool(BaseTool):
 
     Inherits from BaseTool.
     """
-    def __init__(self, config: Optional[Dict[str, Any]] = None):
+
+    # Configuration schema
+    class Config(BaseSettings):
+        """Configuration for the research tool
+        
+        Automatically reads from environment variables with RESEARCH_TOOL_ prefix.
+        Example: RESEARCH_TOOL_SPACY_MODEL -> spacy_model
+        """
+
+        model_config = SettingsConfigDict(env_prefix="RESEARCH_TOOL_")
+
+        max_workers: int = Field(
+            default=min(32, (os.cpu_count() or 4) * 2),
+            description="Maximum number of worker threads",
+        )
+        spacy_model: str = Field(default="en_core_web_sm", description="Default spaCy model to use")
+        max_text_length: int = Field(default=10_000, description="Maximum text length for inputs")
+        allowed_spacy_models: List[str] = Field(
+            default=["en_core_web_sm", "zh_core_web_sm"],
+            description="Allowed spaCy models",
+        )
+
+    # Schema definitions
+    class Mill_agreementSchema(BaseModel):
+        """Schema for mill_agreement operation"""
+
+        cases: List[Dict[str, Any]] = Field(description="List of cases with attributes and outcomes. Each case should have 'attrs' (dict of attributes) and 'outcome' (boolean)")
+
+    class Mill_differenceSchema(BaseModel):
+        """Schema for mill_difference operation"""
+
+        positive_case: Dict[str, Any] = Field(description="Positive case with attributes and outcome. Should have 'attrs' (dict of attributes) and 'outcome' (boolean)")
+        negative_case: Dict[str, Any] = Field(description="Negative case with attributes and outcome. Should have 'attrs' (dict of attributes) and 'outcome' (boolean)")
+
+    class Mill_jointSchema(BaseModel):
+        """Schema for mill_joint operation"""
+
+        positive_cases: List[Dict[str, Any]] = Field(description="List of positive cases. Each case should have 'attrs' (dict of attributes) and 'outcome' (boolean)")
+        negative_cases: List[Dict[str, Any]] = Field(description="List of negative cases. Each case should have 'attrs' (dict of attributes) and 'outcome' (boolean)")
+
+    class Mill_residuesSchema(BaseModel):
+        """Schema for mill_residues operation"""
+
+        cases: List[Dict[str, Any]] = Field(description="List of cases with attributes and effects. Each case should have 'attrs' (dict of attributes) and 'effects' (list of effect names)")
+        known_causes: Dict[str, List[str]] = Field(description="Dictionary mapping effect names to lists of known cause attribute names")
+
+    class Mill_concomitantSchema(BaseModel):
+        """Schema for mill_concomitant operation"""
+
+        cases: List[Dict[str, Any]] = Field(description="List of cases with attributes. Each case should have 'attrs' (dict of attributes with numeric values)")
+        factor: str = Field(description="Name of the factor attribute to analyze")
+        effect: str = Field(description="Name of the effect attribute to analyze")
+
+    class InductionSchema(BaseModel):
+        """Schema for induction operation"""
+
+        examples: List[str] = Field(description="List of example text strings to generalize patterns from")
+        max_keywords: int = Field(default=10, description="Maximum number of keywords/patterns to extract from the examples")
+
+    class DeductionSchema(BaseModel):
+        """Schema for deduction operation"""
+
+        premises: List[str] = Field(description="List of premise statement strings to validate against")
+        conclusion: Optional[str] = Field(default=None, description="Optional conclusion statement string to validate. If None, validation will fail")
+
+    class SummarizeSchema(BaseModel):
+        """Schema for summarize operation"""
+
+        text: str = Field(description="Text string to summarize")
+        max_length: int = Field(default=150, description="Maximum length of the summary in words")
+        language: Optional[str] = Field(default=None, description="Optional language code for the text. If None, uses the default spaCy model language")
+
+    def __init__(self, config: Optional[Dict[str, Any]] = None, **kwargs):
         """
         Initialize ResearchTool with settings and resources.
 
         Args:
-            config (Dict, optional): Configuration overrides for ResearchSettings.
+            config (Dict, optional): Configuration overrides for ResearchTool.
+            **kwargs: Additional arguments passed to BaseTool (e.g., tool_name)
 
         Raises:
             ValueError: If config contains invalid settings.
+
+        Configuration is automatically loaded by BaseTool from:
+        1. Explicit config dict (highest priority)
+        2. YAML config files (config/tools/research.yaml)
+        3. Environment variables (via dotenv from .env files)
+        4. Tool defaults (lowest priority)
         """
-        super().__init__(config)
-        self.settings = ResearchSettings()
-        if config:
-            try:
-                self.settings = self.settings.model_validate({**self.settings.model_dump(), **config})
-            except ValidationError as e:
-                raise ValueError(f"Invalid configuration: {e}")
+        super().__init__(config, **kwargs)
+
+        # Configuration is automatically loaded by BaseTool into self._config_obj
+        # Access config via self._config_obj (BaseSettings instance)
+        self.config = self._config_obj if self._config_obj else self.Config()
+
         self.logger = logging.getLogger(__name__)
         if not self.logger.handlers:
             handler = logging.StreamHandler()
-            handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(message)s'))
+            handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
             self.logger.addHandler(handler)
         self.logger.setLevel(logging.INFO)
         self._spacy_nlp: Optional[Language] = None
 
     def __del__(self):
         """Clean up resources when the object is destroyed."""
-        if hasattr(self, '_spacy_nlp') and self._spacy_nlp is not None:
+        if hasattr(self, "_spacy_nlp") and self._spacy_nlp is not None:
             self._spacy_nlp = None
 
     def _get_spacy(self) -> Language:
@@ -98,9 +157,9 @@ class ResearchTool(BaseTool):
             ResearchToolError: If the spaCy model is invalid.
         """
         if self._spacy_nlp is None:
-            if self.settings.spacy_model not in self.settings.allowed_spacy_models:
-                raise ResearchToolError(f"Invalid spaCy model '{self.settings.spacy_model}', expected {self.settings.allowed_spacy_models}")
-            self._spacy_nlp = spacy.load(self.settings.spacy_model, disable=["textcat"])
+            if self.config.spacy_model not in self.config.allowed_spacy_models:
+                raise ResearchToolError(f"Invalid spaCy model '{self.config.spacy_model}', expected {self.config.allowed_spacy_models}")
+            self._spacy_nlp = spacy.load(self.config.spacy_model, disable=["textcat"])
         return self._spacy_nlp
 
     def mill_agreement(self, cases: List[Dict[str, Any]]) -> Dict[str, Any]:
@@ -117,13 +176,13 @@ class ResearchTool(BaseTool):
             FileOperationError: If processing fails.
         """
         try:
-            truthy = [c['attrs'] for c in cases if c.get('outcome')]
+            truthy = [c["attrs"] for c in cases if c.get("outcome")]
             if not truthy:
-                return {'common_factors': []}
+                return {"common_factors": []}
             common = set(k for k, v in truthy[0].items() if v)
             for attrs in truthy[1:]:
                 common &= set(k for k, v in attrs.items() if v)
-            return {'common_factors': list(common)}
+            return {"common_factors": list(common)}
         except Exception as e:
             raise FileOperationError(f"Failed to process mill_agreement: {str(e)}")
 
@@ -142,14 +201,18 @@ class ResearchTool(BaseTool):
             FileOperationError: If processing fails.
         """
         try:
-            pos = {k for k, v in positive_case.get('attrs', {}).items() if v}
-            neg = {k for k, v in negative_case.get('attrs', {}).items() if v}
+            pos = {k for k, v in positive_case.get("attrs", {}).items() if v}
+            neg = {k for k, v in negative_case.get("attrs", {}).items() if v}
             diff = pos - neg
-            return {'difference_factors': list(diff)}
+            return {"difference_factors": list(diff)}
         except Exception as e:
             raise FileOperationError(f"Failed to process mill_difference: {str(e)}")
 
-    def mill_joint(self, positive_cases: List[Dict[str, Any]], negative_cases: List[Dict[str, Any]]) -> Dict[str, Any]:
+    def mill_joint(
+        self,
+        positive_cases: List[Dict[str, Any]],
+        negative_cases: List[Dict[str, Any]],
+    ) -> Dict[str, Any]:
         """
         Combine Mill's Method of Agreement and Difference to identify causal factors.
 
@@ -164,18 +227,18 @@ class ResearchTool(BaseTool):
             FileOperationError: If processing fails.
         """
         try:
-            truthy = [c['attrs'] for c in positive_cases if c.get('outcome')]
+            truthy = [c["attrs"] for c in positive_cases if c.get("outcome")]
             if not truthy:
-                return {'causal_factors': []}
+                return {"causal_factors": []}
             common = set(k for k, v in truthy[0].items() if v)
             for attrs in truthy[1:]:
                 common &= set(k for k, v in attrs.items() if v)
-            falsy = [c['attrs'] for c in negative_cases if not c.get('outcome')]
+            falsy = [c["attrs"] for c in negative_cases if not c.get("outcome")]
             if not falsy:
-                return {'causal_factors': list(common)}
+                return {"causal_factors": list(common)}
             for attrs in falsy:
                 common -= set(k for k, v in attrs.items() if v)
-            return {'causal_factors': list(common)}
+            return {"causal_factors": list(common)}
         except Exception as e:
             raise FileOperationError(f"Failed to process mill_joint: {str(e)}")
 
@@ -196,15 +259,15 @@ class ResearchTool(BaseTool):
         try:
             residual = {}
             for case in cases:
-                effects = case.get('effects', {})
-                attrs = set(k for k, v in case.get('attrs', {}).items() if v)
+                effects = case.get("effects", {})
+                attrs = set(k for k, v in case.get("attrs", {}).items() if v)
                 for effect in effects:
                     if effect in known_causes:
                         known = set(known_causes[effect])
                         residual[effect] = list(attrs - known)
                     else:
                         residual[effect] = list(attrs)
-            return {'residual_causes': residual}
+            return {"residual_causes": residual}
         except Exception as e:
             raise FileOperationError(f"Failed to process mill_residues: {str(e)}")
 
@@ -224,16 +287,17 @@ class ResearchTool(BaseTool):
             FileOperationError: If processing fails.
         """
         try:
-            factor_vals = [case['attrs'].get(factor, 0) for case in cases]
-            effect_vals = [case['attrs'].get(effect, 0) for case in cases]
+            factor_vals = [case["attrs"].get(factor, 0) for case in cases]
+            effect_vals = [case["attrs"].get(effect, 0) for case in cases]
             if len(factor_vals) < 2:
-                return {'correlation': 0.0, 'pvalue': 1.0}
-            
+                return {"correlation": 0.0, "pvalue": 1.0}
+
             # Convert to numpy arrays to avoid PyTorch compatibility issues
             import numpy as np
+
             factor_array = np.array(factor_vals, dtype=np.float64)
             effect_array = np.array(effect_vals, dtype=np.float64)
-            
+
             # Calculate correlation using numpy if scipy fails
             try:
                 corr, pval = pearsonr(factor_array, effect_array)
@@ -241,17 +305,19 @@ class ResearchTool(BaseTool):
                 # Fallback to numpy correlation calculation
                 self.logger.warning(f"scipy pearsonr failed ({e}), using numpy fallback")
                 corr = np.corrcoef(factor_array, effect_array)[0, 1]
-                # Simple p-value approximation (not statistically rigorous but functional)
+                # Simple p-value approximation (not statistically rigorous but
+                # functional)
                 n = len(factor_array)
                 if n <= 2:
                     pval = 1.0
                 else:
                     # Approximate p-value using t-distribution
                     t_stat = corr * np.sqrt((n - 2) / (1 - corr**2 + 1e-10))
-                    from scipy.stats import t
+                    from scipy.stats import t  # type: ignore[import-untyped]
+
                     pval = 2 * (1 - t.cdf(abs(t_stat), n - 2))
-            
-            return {'correlation': float(corr), 'pvalue': float(pval)}
+
+            return {"correlation": float(corr), "pvalue": float(pval)}
         except Exception as e:
             raise FileOperationError(f"Failed to process mill_concomitant: {str(e)}")
 
@@ -275,10 +341,10 @@ class ResearchTool(BaseTool):
             patterns = []
             for doc in docs:
                 patterns.extend([chunk.text.lower() for chunk in doc.noun_chunks])
-                patterns.extend([token.lemma_.lower() for token in doc if token.pos_ == 'VERB'])
+                patterns.extend([token.lemma_.lower() for token in doc if token.pos_ == "VERB"])
             counter = Counter(patterns)
             common = [word for word, count in counter.most_common() if count > 1][:max_keywords]
-            return {'patterns': common}
+            return {"patterns": common}
         except Exception as e:
             raise FileOperationError(f"Failed to process induction: {str(e)}")
 
@@ -301,14 +367,18 @@ class ResearchTool(BaseTool):
             premises_docs = [nlp(p) for p in premises]
             conclusion_doc = nlp(conclusion) if conclusion else None
             if not conclusion_doc:
-                return {'valid': False, 'conclusion': None, 'reason': 'No conclusion provided'}
-            premise_entities = set()
-            premise_predicates = set()
+                return {
+                    "valid": False,
+                    "conclusion": None,
+                    "reason": "No conclusion provided",
+                }
+            premise_entities: Set[str] = set()
+            premise_predicates: Set[str] = set()
             for doc in premises_docs:
                 premise_entities.update(ent.text.lower() for ent in doc.ents)
-                premise_predicates.update(token.lemma_.lower() for token in doc if token.pos_ == 'VERB')
+                premise_predicates.update(token.lemma_.lower() for token in doc if token.pos_ == "VERB")
             conclusion_entities = set(ent.text.lower() for ent in conclusion_doc.ents)
-            conclusion_predicates = set(token.lemma_.lower() for token in conclusion_doc if token.pos_ == 'VERB')
+            conclusion_predicates = set(token.lemma_.lower() for token in conclusion_doc if token.pos_ == "VERB")
             entities_valid = conclusion_entities.issubset(premise_entities)
             predicates_valid = conclusion_predicates.issubset(premise_predicates)
             valid = entities_valid and predicates_valid
@@ -316,10 +386,10 @@ class ResearchTool(BaseTool):
                 "Conclusion matches premise patterns."
                 if valid
                 else f"Conclusion contains unmatched {'entities' if not entities_valid else ''} "
-                     f"{'and ' if not entities_valid and not predicates_valid else ''}"
-                     f"{'predicates' if not predicates_valid else ''}."
+                f"{'and ' if not entities_valid and not predicates_valid else ''}"
+                f"{'predicates' if not predicates_valid else ''}."
             )
-            return {'valid': valid, 'conclusion': conclusion, 'reason': reason}
+            return {"valid": valid, "conclusion": conclusion, "reason": reason}
         except Exception as e:
             raise FileOperationError(f"Failed to process deduction: {str(e)}")
 
@@ -344,20 +414,20 @@ class ResearchTool(BaseTool):
             sentences = [sent.text for sent in doc.sents]
             if not sentences:
                 return ""
-            keywords = [token.lemma_.lower() for token in doc if token.pos_ in ('NOUN', 'VERB', 'ADJ') and not token.is_stop]
+            keywords = [token.lemma_.lower() for token in doc if token.pos_ in ("NOUN", "VERB", "ADJ") and not token.is_stop]
             keyword_freq = Counter(keywords)
             scores = []
             for sent in sentences:
                 sent_doc = nlp(sent)
-                sent_keywords = [token.lemma_.lower() for token in sent_doc if token.pos_ in ('NOUN', 'VERB', 'ADJ')]
+                sent_keywords = [token.lemma_.lower() for token in sent_doc if token.pos_ in ("NOUN", "VERB", "ADJ")]
                 score = sum(keyword_freq.get(k, 0) for k in sent_keywords) / (len(sent_keywords) + 1)
                 scores.append((sent, score))
             scores.sort(key=lambda x: x[1], reverse=True)
-            selected = [sent for sent, _ in scores[:max(1, max_length // 50)]]
-            summary = ' '.join(selected)
+            selected = [sent for sent, _ in scores[: max(1, max_length // 50)]]
+            summary = " ".join(selected)
             words = summary.split()
             if len(words) > max_length:
-                summary = ' '.join(words[:max_length]) + '...'
+                summary = " ".join(words[:max_length]) + "..."
             return summary
         except Exception as e:
             raise FileOperationError(f"Failed to process summarize: {str(e)}")