import asyncio
import gradio as gr
from langchain_community.llms import HuggingFaceEndpoint
from langchain_core.prompts import PromptTemplate
from langchain.agents import create_react_agent, AgentExecutor
from langchain_core.tools import BaseTool
from typing import List
import yaml
import os
import json
import logging
from functools import lru_cache
import time
import pygments
from pygments.lexers import get_lexer_by_name
from pygments.formatters import HtmlFormatter
# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
def load_config():
config_path = os.path.join(os.path.dirname(__file__), "config.yaml")
try:
with open(config_path, "r") as config_file:
return yaml.safe_load(config_file)
except FileNotFoundError:
logger.warning("Config file not found. Using default configuration.")
return {
"model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
"hf_api_token": os.environ.get("HUGGINGFACEHUB_API_TOKEN", "your_default_token_here")
}
config = load_config()
class AgentInitializationError(Exception):
pass
class CodeGenerationTool(BaseTool):
name = "CodeGeneration"
description = "Generates code based on a prompt"
@lru_cache(maxsize=100)
def _run(self, prompt: str) -> str:
logger.info(f"Generating code for prompt: {prompt}")
if "Flask app structure" in prompt:
return self.generate_flask_app_structure()
elif "binary search algorithm" in prompt:
return self.generate_binary_search()
elif "responsive navbar" in prompt:
return self.generate_responsive_navbar()
else:
return f"Generated code placeholder for: {prompt}"
async def _arun(self, prompt: str) -> str:
return self._run(prompt)
def generate_flask_app_structure(self):
return """
# app.py
from flask import Flask, render_template
app = Flask(__name__)
@app.route('/')
def home():
return render_template('index.html')
if __name__ == '__main__':
app.run(debug=True)
# templates/index.html
Flask App
Welcome to Flask!
"""
def generate_binary_search(self):
return """
def binary_search(arr, target):
left, right = 0, len(arr) - 1
while left <= right:
mid = (left + right) // 2
if arr[mid] == target:
return mid
elif arr[mid] < target:
left = mid + 1
else:
right = mid - 1
return -1 # Target not found
# Example usage
sorted_array = [1, 3, 5, 7, 9, 11, 13, 15]
target = 7
result = binary_search(sorted_array, target)
print(f"Target {target} found at index: {result}")
"""
def generate_responsive_navbar(self):
return """
"""
class Agent:
def __init__(self, name: str, description: str, tools: List[BaseTool]):
self.name = name
self.description = description
try:
self.llm = HuggingFaceEndpoint(
repo_id=config["model"],
task="text-generation",
model_kwargs={"temperature": 0.7, "max_length": 1024},
huggingfacehub_api_token=config["hf_api_token"]
)
self.prompt_template = PromptTemplate(
template="You are {name}, {description}. Respond to the following: {input}",
input_variables=["name", "description", "input"]
)
self.agent = create_react_agent(self.llm, tools, self.prompt_template)
self.agent_executor = AgentExecutor(agent=self.agent, tools=tools, verbose=True)
except Exception as e:
logger.error(f"Failed to initialize agent: {e}")
raise AgentInitializationError(f"Failed to initialize agent: {e}")
async def run(self, input_text: str) -> str:
try:
result = await self.agent_executor.arun(input_text)
return result
except Exception as e:
logger.error(f"Error in agent execution: {e}")
return f"Error: {str(e)}"
class CodeFusion:
def __init__(self):
code_gen_tool = CodeGenerationTool()
self.agents = [
Agent("CodeFusion_Structure", "App Structure Designer", [code_gen_tool]),
Agent("CodeFusion_Logic", "Logic Implementation Expert", [code_gen_tool]),
Agent("CodeFusion_UI", "User Interface Designer", [code_gen_tool])
]
async def run(self, input_text: str) -> str:
results = []
for agent in self.agents:
result = await agent.run(input_text)
results.append(f"{agent.name}: {result}")
return "\n\n".join(results)
code_fusion = CodeFusion()
def highlight_code(code: str, language: str) -> str:
lexer = get_lexer_by_name(language, stripall=True)
formatter = HtmlFormatter(style="monokai")
return pygments.highlight(code, lexer, formatter)
def save_code_to_file(code: str, filename: str) -> str:
try:
with open(filename, 'w') as f:
f.write(code)
return f"Code saved to {filename}"
except Exception as e:
logger.error(f"Error saving code to file: {e}")
return f"Error saving code: {str(e)}"
async def chat(message, history):
start_time = time.time()
response = await code_fusion.run(message)
end_time = time.time()
# Highlight code in the response
highlighted_response = response
for lang in ['python', 'html', 'css', 'javascript']:
if f"```{lang}" in response:
code = response.split(f"```{lang}")[1].split("```")[0]
highlighted_code = highlight_code(code, lang)
highlighted_response = highlighted_response.replace(f"```{lang}{code}```", highlighted_code)
# Save code to file if requested
if "save code" in message.lower():
filename = f"generated_code_{int(time.time())}.py"
save_result = save_code_to_file(response, filename)
highlighted_response += f"\n\n{save_result}"
execution_time = end_time - start_time
highlighted_response += f"\n\nExecution time: {execution_time:.2f} seconds"
return highlighted_response
async def main():
iface = gr.ChatInterface(
fn=chat,
title="CodeFusion AI",
description="Your AI-powered coding assistant",
examples=[
"Create a basic Flask app structure",
"Implement a binary search algorithm in Python",
"Design a responsive navbar using HTML and CSS",
"Generate a simple REST API using Flask and save the code"
],
retry_btn=None,
undo_btn="Delete Previous",
clear_btn="Clear",
)
await iface.launch()
if __name__ == "__main__":
asyncio.run(main())