app.py

import streamlit as st
from langchain.agents import create_react_agent, AgentExecutor
from langchain.prompts import PromptTemplate
from langchain_community.llms import HuggingFaceHub
from langchain_huggingface import HuggingFaceEndpoint
from langchain.tools import Tool
from langchain.chains import LLMChain
from typing import List, Dict, Any, Optional

# Base Agent class
class Agent:
    def __init__(self, name, role, tools, knowledge_base=None):
        self.name = name
        self.role = role
        self.tools = tools
        self.knowledge_base = knowledge_base
        self.llm = HuggingFaceHub(repo_id="google/flan-t5-xl", model_kwargs={"temperature": 0.5})

    def act(self, task, context):
        # This method should be implemented in subclasses
        raise NotImplementedError

# Base Tool class
class BaseTool(Tool):
    def __init__(self, name, description):
        super().__init__(name=name, description=description, func=self.run)
        self.llm = HuggingFaceHub(repo_id="google/flan-t5-xl", model_kwargs={"temperature": 0.5})

    def run(self, arguments: str) -> str:
        raise NotImplementedError
        
# Specific tool implementations
class CodeGenerationTool(BaseTool):
    def __init__(self):
        super().__init__("Code Generation", "Generates code snippets in various languages.")
        self.prompt_template = PromptTemplate(
            input_variables=["language", "task"],
            template="Generate {language} code for: {task}"
        )
        self.chain = LLMChain(llm=self.llm, prompt=self.prompt_template)

    def run(self, arguments: str) -> str:
        language, task = arguments.split(", ", 1)
        return self.chain.run(language=language, task=task)

class DataRetrievalTool(BaseTool):
    def __init__(self):
        super().__init__("Data Retrieval", "Accesses data from APIs, databases, or files.")
        self.prompt_template = PromptTemplate(
            input_variables=["source", "query"],
            template="Retrieve data from {source} based on: {query}"
        )
        self.chain = LLMChain(llm=self.llm, prompt=self.prompt_template)

    def run(self, arguments: str) -> str:
        source, query = arguments.split(", ", 1)
        return self.chain.run(source=source, query=query)

class TextGenerationTool(BaseTool):
    def __init__(self):
        super().__init__("Text Generation", "Generates human-like text based on a given prompt.")
        self.prompt_template = PromptTemplate(
            input_variables=["prompt"],
            template="Generate text based on: {prompt}"
        )
        self.chain = LLMChain(llm=self.llm, prompt=self.prompt_template)

    def run(self, arguments: str) -> str:
        return self.chain.run(prompt=arguments)
        
# Specialized Agent Definitions
class SpecializedAgent(Agent):
    def __init__(self, name, role, tools, knowledge_base=None):
        super().__init__(name, role, tools, knowledge_base)
        self.prompt = PromptTemplate.from_template(
            "You are a specialized AI assistant named {name} with the role of {role}. "
            "Use the following tools to complete your task: {tools}\n\n"
            "Task: {input}\n"
            "Thought: Let's approach this step-by-step:\n"
            "{agent_scratchpad}"
        )
        self.react_agent = create_react_agent(self.llm, self.tools, self.prompt)
        self.agent_executor = AgentExecutor(
            agent=self.react_agent, 
            tools=self.tools, 
            verbose=True,
            max_iterations=5
        )

    def act(self, task, context):
        return self.agent_executor.run(input=task)

class RequirementsAgent(SpecializedAgent):
    def __init__(self):
        super().__init__("RequirementsAnalyst", "Analyzing and refining project requirements", [TextGenerationTool()])

class ArchitectureAgent(SpecializedAgent):
    def __init__(self):
        super().__init__("SystemArchitect", "Designing system architecture", [TextGenerationTool()])

class FrontendAgent(SpecializedAgent):
    def __init__(self):
        super().__init__("FrontendDeveloper", "Developing the frontend", [CodeGenerationTool()])

class BackendAgent(SpecializedAgent):
    def __init__(self):
        super().__init__("BackendDeveloper", "Developing the backend", [CodeGenerationTool(), DataRetrievalTool()])

class DatabaseAgent(SpecializedAgent):
    def __init__(self):
        super().__init__("DatabaseEngineer", "Designing and implementing the database", [CodeGenerationTool(), DataRetrievalTool()])

class TestingAgent(SpecializedAgent):
    def __init__(self):
        super().__init__("QAEngineer", "Creating and executing test plans", [CodeGenerationTool(), TextGenerationTool()])

class DeploymentAgent(SpecializedAgent):
    def __init__(self):
        super().__init__("DevOpsEngineer", "Handling deployment and infrastructure", [CodeGenerationTool(), TextGenerationTool()])

# --- Application Building Sequence ---
class ApplicationBuilder:
    def __init__(self):
        self.agents = [
            RequirementsAgent(),
            ArchitectureAgent(),
            FrontendAgent(),
            BackendAgent(),
            DatabaseAgent(),
            TestingAgent(),
            DeploymentAgent()
        ]
        self.project_state = {}

    def build_application(self, project_description):
        st.write("Starting application building process...")
        for agent in self.agents:
            st.write(f"\n--- {agent.name}'s Turn ---")
            task = self.get_task_for_agent(agent, project_description)
            response = agent.act(task, self.project_state)
            self.project_state[agent.role] = response
            st.write(f"{agent.name}'s output:")
            st.write(response)
        st.write("\nApplication building process completed!")

    def get_task_for_agent(self, agent, project_description):
        tasks = (
            (RequirementsAgent, f"Analyze and refine the requirements for this project: {project_description}"),
            (ArchitectureAgent, f"Design the system architecture based on these requirements: {self.project_state.get('Analyzing and refining project requirements', '')}"),
            (FrontendAgent, f"Develop the frontend based on this architecture: {self.project_state.get('Designing system architecture', '')}"),
            (BackendAgent, f"Develop the backend based on this architecture: {self.project_state.get('Designing system architecture', '')}"),
            (DatabaseAgent, f"Design and implement the database based on this architecture: {self.project_state.get('Designing system architecture', '')}"),
            (TestingAgent, f"Create a test plan for this application: {project_description}"),
            (DeploymentAgent, f"Create a deployment plan for this application: {project_description}")
        )
        
        for agent_class, task in tasks:
            if isinstance(agent, agent_class):
                return task
        
        return f"Contribute to the project: {project_description}"

# --- Streamlit App ---
st.title("CODEFUSSION ☄ - Full-Stack Application Builder")

project_description = st.text_area("Enter your project description:")

if st.button("Build Application"):
    if project_description:
        app_builder = ApplicationBuilder()
        app_builder.build_application(project_description)
    else:
        st.write("Please enter a project description.")

# Display information about the agents
st.sidebar.title("Agent Information")
app_builder = ApplicationBuilder()
for agent in app_builder.agents:
    st.sidebar.write(f"--- {agent.name} ---")
    st.sidebar.write(f"Role: {agent.role}")
    st.sidebar.write("Tools:")
    for tool in agent.tools:
        st.sidebar.write(f"- {tool.name}")
    st.sidebar.write("")