Update app.py

app.py

@@ -1,10 +1,142 @@
-# app.py
+import os
 import sys
+import logging
+from typing import Optional
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler(sys.stdout)]
+)
+logger = logging.getLogger(__name__)
+
+def setup_environment():
+    """Setup environment variables and cache directories"""
+    try:
+        cache_dir = "/tmp/hf"
+        os.environ["TRANSFORMERS_CACHE"] = cache_dir
+        os.environ["HF_HOME"] = cache_dir
+        
+        # Create cache directory if it doesn't exist
+        os.makedirs(cache_dir, exist_ok=True)
+        logger.info(f"Cache directory set up at {cache_dir}")
+        return True
+    except Exception as e:
+        logger.error(f"Failed to setup environment: {str(e)}")
+        return False
+
+def initialize_model():
+    """Initialize the model and pipeline with error handling"""
+    try:
+        from transformers import pipeline
+        logger.info("Loading model pipeline...")
+        
+        pipe = pipeline(
+            "text-generation",
+            model="bigscience/bloom-560m",  # Using a smaller model
+            tokenizer="bigscience/bloom-560m",
+            device=-1,  # Force CPU
+            max_new_tokens=256,
+            force_download=True,
+            low_cpu_mem_usage=True,  # Added for better memory management
+        )
+        logger.info("Model pipeline loaded successfully")
+        return pipe
+    except Exception as e:
+        logger.error(f"Failed to initialize model: {str(e)}")
+        return None
+
+def setup_agent(pipe):
+    """Setup LangChain agent with the model pipeline"""
+    try:
+        from langchain_community.llms import HuggingFacePipeline
+        from langchain.agents import initialize_agent
+        from langchain.agents.agent_types import AgentType
+        from langchain.tools import BaseTool
+
+        # Initialize LLM
+        llm = HuggingFacePipeline(pipeline=pipe)
+
+        # Define the tool
+        class MachineryReportTool(BaseTool):
+            name = "machinery_report"
+            description = (
+                "Generates a report on mini construction equipment including "
+                "specifications and market analysis."
+            )
+
+            def _run(self, query: str) -> str:
+                # Simplified report for memory efficiency
+                return """
+                Mini Construction Equipment Report:
+                1. Basic Gas-Powered Unit: $3,700, 14HP
+                2. DRT 450: $5,500-6,300, Honda motor
+                3. Mini Skid Steer: $15,000, Kubota diesel
+                Market Analysis: 30-40% cheaper than US equivalents
+                """
+
+            def _arun(self, query: str) -> str:
+                raise NotImplementedError("Async not supported.")
+
+        # Initialize agent with minimal settings
+        agent = initialize_agent(
+            tools=[MachineryReportTool()],
+            llm=llm,
+            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+            verbose=True,
+            max_iterations=1,
+            early_stopping_method="generate"  # Added to prevent infinite loops
+        )
+        logger.info("Agent initialized successfully")
+        return agent
+    except Exception as e:
+        logger.error(f"Failed to setup agent: {str(e)}")
+        return None
+
+def run_agent_query(agent, query: str) -> Optional[str]:
+    """Run a query through the agent with error handling"""
+    try:
+        logger.info("Running agent query...")
+        result = agent.run(query)
+        logger.info("Query completed successfully")
+        return result
+    except Exception as e:
+        logger.error(f"Error running agent query: {str(e)}")
+        return f"Error processing query: {str(e)}"
 
 def main():
-    print("Hello from my Python script on Hugging Face Spaces!")
-    # Your code here
-    # e.g., run your LangChain script, do computations, etc.
+    """Main function to run the agent"""
+    logger.info("Starting application...")
+    
+    # Setup environment
+    if not setup_environment():
+        return
+    
+    # Initialize model
+    pipe = initialize_model()
+    if pipe is None:
+        return
+    
+    # Setup agent
+    agent = setup_agent(pipe)
+    if agent is None:
+        return
+    
+    # Define query
+    query = (
+        "Please provide a brief overview of the mini construction equipment "
+        "project, focusing on key specifications and market positioning."
+    )
+    
+    # Run query
+    result = run_agent_query(agent, query)
+    
+    # Print result
+    if result:
+        print("\n===== AGENT OUTPUT =====")
+        print(result)
+        print("=======================")
 
 if __name__ == "__main__":
     main()