Fixed token length issues and optimized prompts

app.py

@@ -59,10 +59,15 @@ def get_llm():
                 "text-generation",
                 model=model,
                 tokenizer=tokenizer,
-                max_length=500,
+                max_new_tokens=500,  
+                pad_token_id=tokenizer.eos_token_id,
                 temperature=0.7,
+                return_full_text=False,  
+            )
+            llm = HuggingFacePipeline(
+                pipeline=hf_pipeline,
+                model_kwargs={"max_length": 2048}  
             )
-            llm = HuggingFacePipeline(pipeline=hf_pipeline)
         return llm
     except Exception as e:
         logger.error(f"Failed to get LLM: {str(e)}")
@@ -809,130 +814,68 @@ class EnhancedAIFlow(AIFlow):
             context = self.contexts[role]
             dependencies_output = []
             
-            # Gather outputs from dependencies
+            # Gather outputs from dependencies (limited to last 1000 chars)
             for dep_role in context.dependencies:
                 dep_context = self.contexts[dep_role]
                 if dep_context.state == FlowState.COMPLETED and "output" in dep_context.artifacts:
-                    dependencies_output.append(f"## {dep_role.value} Output:\n{dep_context.artifacts['output']}")
+                    output = dep_context.artifacts['output']
+                    if len(output) > 1000:  
+                        output = output[:997] + "..."
+                    dependencies_output.append(f"## {dep_role.value} Output:\n{output}")
             
-            # Build role-specific prompts
+            # Build role-specific prompts (with size limits)
             role_prompts = {
-                AgentRole.ARCHITECT: """You are a Software Architect. Design the high-level architecture for:
-Requirements:
-{requirements}
-
-Focus on:
-1. Overall system design
-2. Component interactions
-3. Technology stack
-4. Data flow
-5. Scalability considerations""",
+                AgentRole.ARCHITECT: """Design the high-level architecture (brief overview):
+Requirements: {requirements}
+Focus: system design, components, tech stack, data flow, scalability""",
                 
-                AgentRole.UI_DESIGNER: """You are a UI/UX Designer. Design the user interface for:
-Requirements:
-{requirements}
-
-Previous Designs:
-{dependencies}
-
-Focus on:
-1. User experience
-2. Component layout
-3. Responsive design
-4. Visual hierarchy
-5. Accessibility""",
+                AgentRole.UI_DESIGNER: """Design the UI (key elements):
+Requirements: {requirements}
+Previous: {dependencies}
+Focus: UX, layout, responsiveness, themes""",
                 
-                AgentRole.BACKEND_DEVELOPER: """You are a Backend Developer. Implement the server-side logic for:
-Requirements:
-{requirements}
-
-Architecture:
-{dependencies}
-
-Focus on:
-1. API design
-2. Business logic
-3. Data validation
-4. Error handling
-5. Performance""",
+                AgentRole.BACKEND_DEVELOPER: """Implement core backend logic:
+Requirements: {requirements}
+Architecture: {dependencies}
+Focus: API, business logic, validation""",
                 
-                AgentRole.DATABASE_ENGINEER: """You are a Database Engineer. Design the data layer for:
-Requirements:
-{requirements}
-
-System Context:
-{dependencies}
-
-Focus on:
-1. Schema design
-2. Data relationships
-3. Indexing strategy
-4. Query optimization
-5. Data integrity""",
+                AgentRole.DATABASE_ENGINEER: """Design data layer:
+Requirements: {requirements}
+Context: {dependencies}
+Focus: schema, relationships, optimization""",
                 
-                AgentRole.SECURITY_EXPERT: """You are a Security Expert. Review and enhance security for:
-Requirements:
-{requirements}
-
-System Context:
-{dependencies}
-
-Focus on:
-1. Authentication
-2. Authorization
-3. Data protection
-4. Security best practices
-5. Compliance requirements""",
+                AgentRole.SECURITY_EXPERT: """Review security:
+Requirements: {requirements}
+Context: {dependencies}
+Focus: auth, data protection, best practices""",
                 
-                AgentRole.QA_ENGINEER: """You are a QA Engineer. Design the testing strategy for:
-Requirements:
-{requirements}
-
-Implementation Details:
-{dependencies}
-
-Focus on:
-1. Test coverage
-2. Test automation
-3. Edge cases
-4. Performance testing
-5. Security testing""",
+                AgentRole.QA_ENGINEER: """Design testing:
+Requirements: {requirements}
+Implementation: {dependencies}
+Focus: coverage, automation, edge cases""",
                 
-                AgentRole.DEVOPS_ENGINEER: """You are a DevOps Engineer. Set up the deployment pipeline for:
-Requirements:
-{requirements}
-
-System Context:
-{dependencies}
-
-Focus on:
-1. CI/CD pipeline
-2. Infrastructure as code
-3. Monitoring
-4. Scaling
-5. Disaster recovery""",
+                AgentRole.DEVOPS_ENGINEER: """Setup deployment:
+Requirements: {requirements}
+Context: {dependencies}
+Focus: CI/CD, infrastructure, monitoring""",
                 
-                AgentRole.DOCUMENTATION_WRITER: """You are a Technical Writer. Create comprehensive documentation for:
-Requirements:
-{requirements}
-
-System Details:
-{dependencies}
-
-Focus on:
-1. Setup instructions
-2. API documentation
-3. User guides
-4. Architecture overview
-5. Troubleshooting guides"""
+                AgentRole.DOCUMENTATION_WRITER: """Create docs:
+Requirements: {requirements}
+System: {dependencies}
+Focus: setup, API docs, guides"""
             }
             
             # Get the base prompt for the role
             base_prompt = role_prompts.get(role, "")
             
+            # Truncate requirements if too long
+            requirements = self.requirements
+            if len(requirements) > 1000:
+                requirements = requirements[:997] + "..."
+            
             # Format the prompt with requirements and dependencies
             formatted_prompt = base_prompt.format(
-                requirements=self.requirements,
+                requirements=requirements,
                 dependencies="\n\n".join(dependencies_output) if dependencies_output else "No previous context available."
             )