Update app.py

app.py

@@ -1,7 +1,8 @@
 import os
 import sys
 import logging
-from typing import Optional
+from typing import Optional, Any
+from typing_extensions import override
 
 # Configure logging
 logging.basicConfig(
@@ -34,12 +35,12 @@ def initialize_model():
         
         pipe = pipeline(
             "text-generation",
-            model="bigscience/bloom-560m",  # Using a smaller model
+            model="bigscience/bloom-560m",
             tokenizer="bigscience/bloom-560m",
-            device=-1,  # Force CPU
+            device=-1,
             max_new_tokens=256,
             force_download=True,
-            low_cpu_mem_usage=True,  # Added for better memory management
+            low_cpu_mem_usage=True,
         )
         logger.info("Model pipeline loaded successfully")
         return pipe
@@ -50,7 +51,7 @@ def initialize_model():
 def setup_agent(pipe):
     """Setup LangChain agent with the model pipeline"""
     try:
-        from langchain_community.llms import HuggingFacePipeline
+        from langchain_huggingface import HuggingFacePipeline  # Updated import
         from langchain.agents import initialize_agent
         from langchain.agents.agent_types import AgentType
         from langchain.tools import BaseTool
@@ -58,16 +59,16 @@ def setup_agent(pipe):
         # Initialize LLM
         llm = HuggingFacePipeline(pipeline=pipe)
 
-        # Define the tool
+        # Define the tool with proper type annotations
         class MachineryReportTool(BaseTool):
-            name = "machinery_report"
-            description = (
+            name: str = "machinery_report"  # Type annotation added
+            description: str = (  # Type annotation added
                 "Generates a report on mini construction equipment including "
                 "specifications and market analysis."
             )
 
+            @override
             def _run(self, query: str) -> str:
-                # Simplified report for memory efficiency
                 return """
                 Mini Construction Equipment Report:
                 1. Basic Gas-Powered Unit: $3,700, 14HP
@@ -76,7 +77,8 @@ def setup_agent(pipe):
                 Market Analysis: 30-40% cheaper than US equivalents
                 """
 
-            def _arun(self, query: str) -> str:
+            @override
+            def _arun(self, query: str) -> Any:
                 raise NotImplementedError("Async not supported.")
 
         # Initialize agent with minimal settings
@@ -86,7 +88,7 @@ def setup_agent(pipe):
             agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
             verbose=True,
             max_iterations=1,
-            early_stopping_method="generate"  # Added to prevent infinite loops
+            early_stopping_method="generate"
         )
         logger.info("Agent initialized successfully")
         return agent