Upload 6 files

agents/product_manager_agent.py

@@ -1,6 +1,8 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 from langchain_core.messages import AIMessage
+import asyncio
+from typing import Generator, Dict, Any
 
 MODEL_REPO = "Rahul-8799/product_manager_mistral"
 
@@ -11,15 +13,43 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 
-def run(state: dict) -> dict:
-    """Generates structured product requirements from user input prompt."""
-    messages = state["messages"]
-    prompt = messages[-1].content
-
+async def stream_inference(prompt: str) -> Generator[str, None, None]:
+    """Stream the model's output token by token"""
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
-    output_ids = model.generate(input_ids, max_new_tokens=3000)
-    output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    
+    # Generate tokens one by one
+    for _ in range(100):  # Limit to 100 tokens for streaming demo
+        output_ids = model.generate(
+            input_ids,
+            max_new_tokens=1,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        
+        # Get the new token
+        new_token = output_ids[0][-1]
+        if new_token == tokenizer.eos_token_id:
+            break
+            
+        # Decode and yield the token
+        token_text = tokenizer.decode([new_token])
+        yield token_text
+        
+        # Update input_ids for next iteration
+        input_ids = output_ids
+        
+        # Small delay to simulate streaming
+        await asyncio.sleep(0.05)
 
+async def run(state: Dict[str, Any]) -> Dict[str, Any]:
+    """Product Manager generates structured product requirements with streaming output"""
+    messages = state["messages"]
+    prompt = messages[-1].content
+    
+    # Stream the output
+    output = ""
+    async for token in stream_inference(prompt):
+        output += token
+    
     return {
         "messages": [AIMessage(content=output)],
         "chat_log": state["chat_log"] + [{"role": "Product Manager", "content": output}],

agents/software_engineer_agent.py

@@ -18,7 +18,7 @@ def run(state: dict) -> dict:
     
     # Enhance the prompt with UI implementation guidelines
     enhanced_prompt = f"""
-    Objective
+Objective
 
 Generate modern, responsive, and accessible UI code that is visually appealing and adheres to current frontend development best practices.
 

agents/ui_designer_agent.py

@@ -2,7 +2,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 from langchain_core.messages import AIMessage
 
-MODEL_REPO = "Rahul-8799/product_manager_mistral"
+MODEL_REPO = "Rahul-8799/ui_designer_mistral"  # You'll need to fine-tune this model
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(