Fixed model loading and agent initialization

app.py

@@ -16,7 +16,7 @@ import networkx as nx
 from langchain.prompts import PromptTemplate
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from langchain_community.llms import HuggingFacePipeline
+from langchain.llms import HuggingFacePipeline
 from langchain.agents import initialize_agent, Tool
 import subprocess
 import asyncio
@@ -26,33 +26,19 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(
 logger = logging.getLogger(__name__)
 
 # Load the LLM and tokenizer
-MODEL_NAME = "unit-mesh/autodev-coder-deepseek-6.7b-finetunes"
-
 def load_model():
+    """Load the model and tokenizer."""
     try:
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-        # Check available device and memory
-        if torch.cuda.is_available():
-            device = "cuda"
-            # Get available GPU memory
-            gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3  # Convert to GB
-            if gpu_memory < 8:  # If less than 8GB available
-                logger.warning("Limited GPU memory available. Using CPU instead.")
-                device = "cpu"
-        else:
-            device = "cpu"
-            logger.info("No GPU detected. Using CPU.")
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        
+        model_name = "gpt2"  # Using a smaller model for testing
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name)
         
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_NAME,
-            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-            device_map="auto" if device == "cuda" else None,
-            low_cpu_mem_usage=True
-        )
         return tokenizer, model
     except Exception as e:
         logger.error(f"Failed to load model: {str(e)}")
-        raise RuntimeError(f"Model initialization failed: {str(e)}")
+        raise
 
 # Initialize models lazily
 tokenizer = None
@@ -61,44 +47,54 @@ hf_pipeline = None
 llm = None
 
 def get_llm():
-    global tokenizer, model, hf_pipeline, llm
-    if llm is None:
-        tokenizer, model = load_model()
-        hf_pipeline = pipeline(
-            "text-generation",
-            model=model,
-            tokenizer=tokenizer,
-            max_length=500,
-            temperature=0.7,
-        )
-        llm = HuggingFacePipeline(pipeline=hf_pipeline)
-    return llm
+    """Get or initialize the language model."""
+    global llm, tokenizer, model, hf_pipeline
+    
+    try:
+        if llm is None:
+            tokenizer, model = load_model()
+            hf_pipeline = pipeline(
+                "text-generation",
+                model=model,
+                tokenizer=tokenizer,
+                max_length=500,
+                temperature=0.7,
+            )
+            llm = HuggingFacePipeline(pipeline=hf_pipeline)
+        return llm
+    except Exception as e:
+        logger.error(f"Failed to get LLM: {str(e)}")
+        raise
 
-# Lazy initialization of agents
 def get_agent(agent_type):
-    llm = get_llm()
-    return initialize_agent(
-        tools=[
-            Tool(
-                name="Code Formatter",
-                func=lambda x: subprocess.run(["black", "-"], input=x.encode(), capture_output=True).stdout.decode(),
-                description="Formats code using Black.",
-            ),
-            Tool(
-                name="API Generator",
-                func=lambda x: json.dumps({"endpoints": {"example": "POST - Example endpoint."}}),
-                description="Generates API details from code.",
-            ),
-            Tool(
-                name="Task Decomposer",
-                func=lambda x: json.dumps({"tasks": ["Design UI", "Develop Backend", "Test App", "Deploy App"]}),
-                description="Breaks down app requirements into smaller tasks.",
-            ),
-        ],
-        llm=llm,
-        agent="zero-shot-react-description",
-        verbose=True,
-    )
+    """Get or initialize an agent with the specified type."""
+    try:
+        llm = get_llm()
+        return initialize_agent(
+            tools=[
+                Tool(
+                    name="Code Formatter",
+                    func=lambda x: subprocess.run(["black", "-"], input=x.encode(), capture_output=True).stdout.decode(),
+                    description="Formats code using Black.",
+                ),
+                Tool(
+                    name="API Generator",
+                    func=lambda x: json.dumps({"endpoints": {"example": "POST - Example endpoint."}}),
+                    description="Generates API details from code.",
+                ),
+                Tool(
+                    name="Task Decomposer",
+                    func=lambda x: json.dumps({"tasks": ["Design UI", "Develop Backend", "Test App", "Deploy App"]}),
+                    description="Breaks down app requirements into smaller tasks.",
+                ),
+            ],
+            llm=llm,
+            agent="zero-shot-react-description",
+            verbose=True,
+        )
+    except Exception as e:
+        logger.error(f"Failed to get agent: {str(e)}")
+        raise
 
 # Enhanced prompt templates with more specific instructions
 ui_designer_prompt = PromptTemplate(