Removed NLTK-related functionality

database/query_processor.py

@@ -13,7 +13,8 @@ class QueryProcessor:
         self.summarization_model = summarization_model
         self.nlp_model = nlp_model
         self.db_service = db_service
-    
+        logger.info("QueryProcessor initialized")
+
     async def process(
         self,
         query: str,
@@ -22,37 +23,31 @@ class QueryProcessor:
         end_date: Optional[str] = None
     ) -> Dict[str, Any]:
         try:
-            # Validate and parse dates
+            # Date handling
             start_dt = self._parse_date(start_date) if start_date else None
             end_dt = self._parse_date(end_date) if end_date else None
             
-            # Get query embedding
+            # Query processing
             query_embedding = self.embedding_model.encode(query).tolist()
-            logger.debug(f"Query embedding generated for: {query[:50]}...")
-
-            # Extract entities safely
-            entities = self._extract_entities_safely(query)
-            logger.debug(f"Extracted entities: {entities}")
-
-            # Semantic search
-            articles = await self._execute_semantic_search(
+            entities = self.nlp_model.extract_entities(query)
+            
+            # Database search
+            articles = await self._execute_search(
                 query_embedding,
                 start_dt,
                 end_dt,
                 topic,
-                entities
+                [ent[0] for ent in entities]
             )
             
             if not articles:
-                logger.info("No articles found matching criteria")
                 return {"message": "No articles found", "articles": []}
 
-            # Process results and generate summary
-            summary_result = self._generate_summary(articles)
-            
+            # Summary generation
+            summary_data = self._generate_summary(articles)
             return {
-                "summary": summary_result["summary"],
-                "key_sentences": summary_result["key_sentences"],
+                "summary": summary_data["summary"],
+                "key_sentences": summary_data["key_sentences"],
                 "articles": articles,
                 "entities": entities
             }

models/nlp.py

@@ -1,22 +1,35 @@
 import spacy
-import nltk
+from typing import List, Union
+import logging
+
+logger = logging.getLogger(__name__)
 
 class NLPModel:
     def __init__(self):
-        self.nlp = spacy.load("pt_core_news_md")
-        nltk.download('punkt')
-        
-    def __call__(self, text: str):
-        """Makes the model callable like model(text)."""
-        return self.extract_entities(text)  # or another default method
-    
-    def extract_entities(self, text: str):
-        """Ensure this always takes a string and returns entities"""
-        if isinstance(text, list):  # If accidentally passed a list
-            text = " ".join(text)  # Combine into single string
-        doc = self.nlp(text)
-        return [(ent.text.lower(), ent.label_) for ent in doc.ents]
-    
-    
-    def tokenize_sentences(self, text: str):
-        return nltk.sent_tokenize(text)
+        try:
+            # Load spaCy model only
+            self.nlp = spacy.load("pt_core_news_md")
+            logger.info("spaCy model initialized successfully")
+        except Exception as e:
+            logger.error(f"Failed to initialize spaCy model: {str(e)}")
+            raise
+
+    def extract_entities(self, text: Union[str, List[str]]) -> List[tuple]:
+        """Entity extraction using spaCy"""
+        try:
+            if isinstance(text, list):
+                text = " ".join(text)
+            doc = self.nlp(text)
+            return [(ent.text.lower(), ent.label_) for ent in doc.ents]
+        except Exception as e:
+            logger.error(f"Entity extraction failed: {str(e)}")
+            return []
+
+    def tokenize_sentences(self, text: str) -> List[str]:
+        """Sentence tokenization using spaCy"""
+        try:
+            doc = self.nlp(text)
+            return [sent.text for sent in doc.sents]
+        except Exception as e:
+            logger.error(f"Sentence tokenization failed: {str(e)}")
+            return [text]  # Fallback to returning whole text