triAGI-Coder / app.py
acecalisto3's picture
Update app.py
13a7a5d verified
raw
history blame
2.22 kB
import streamlit as st
from huggingface_hub import HfApi
import os
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
from typing import List, Dict
from supplemental import Agent, Tool, CodeGenerationTool, DataRetrievalTool, TextGenerationTool, CodeExecutionTool, CodeDebuggingTool, CodeSummarizationTool, CodeTranslationTool, CodeOptimizationTool, CodeDocumentationTool, ImageGenerationTool, ImageEditingTool, ImageAnalysisTool, Workflow, EnhancedAIAgent
st.title("CODEFUSSION ☄")
# Access Hugging Face API key from secrets
hf_token = st.secrets["hf_token"]
if not hf_token:
st.error("Hugging Face API key not found. Please make sure it is set in the secrets.")
# --- Agent Pool ---
agent_pool = {
"IdeaIntake": EnhancedAIAgent("IdeaIntake", "Idea Intake", ["Data Retrieval", "Code Generation", "Text Generation"], "bigcode/starcoder"),
"CodeBuilder": EnhancedAIAgent("CodeBuilder", "Code Builder", ["Code Generation", "Code Debugging", "Code Optimization"], "bigcode/starcoder"),
"ImageCreator": EnhancedAIAgent("ImageCreator", "Image Creator", ["Image Generation", "Image Editing"], "bigcode/starcoder"),
}
# --- Workflow Definitions ---
class Workflow:
def __init__(self, name, agents, task, description):
self.name = name
self.agents = agents
self.task = task
self.description = description
def run(self, prompt, context):
# Workflow execution logic
for agent in self.agents:
action = agent.act(prompt, context)
# Execute the tool
if action.get("tool"):
tool = next((t for t in agent.tools if t.name == action["tool"]), None)
if tool:
output = tool.run(action["arguments"])
# Update context
context.update(output)
# Example usage
workflow = Workflow(
name="Example Workflow",
agents=[agent_pool["IdeaIntake"], agent_pool["CodeBuilder"]],
task="Generate and debug code",
description="A workflow to generate and debug code using AI agents."
)
context = {}
prompt = "Create a Python function to add two numbers."
workflow.run(prompt, context)
st.write(context)