import gradio as gr import datetime import pandas as pd from groq import Groq from sentence_transformers import SentenceTransformer import chromadb from chromadb.config import Settings import hashlib from typing import TypedDict, Optional, List from langgraph.graph import StateGraph, END import json import tempfile import subprocess import os #from google.colab import userdata #api_key_coder =userdata.get('coder') api_key_coder= os.environ.get('api_key_coder') # --------------------------- # 1. Define State # --------------------------- class CodeAssistantState(TypedDict): user_input: str similar_examples: Optional[List[str]] generated_code: Optional[str] error: Optional[str] task_type: Optional[str] # "generate" or "explain" evaluation_result: Optional[str] # --------------------------- # 2. Initialize Components # --------------------------- # Load data df = pd.read_parquet("hf://datasets/openai/openai_humaneval/openai_humaneval/test-00000-of-00001.parquet") extracted_data = df[['task_id', 'prompt', 'canonical_solution']] # Initialize models and DB embedding_model = SentenceTransformer("all-MiniLM-L6-v2") groq_client = Groq(api_key=api_key_coder) # استبدل بمفتاح API الفعلي client = chromadb.Client(Settings( anonymized_telemetry=False, persist_directory="rag_db" )) collection = client.get_or_create_collection( name="code_examples", metadata={"hnsw:space": "cosine"} ) # --------------------------- # 3. Define Nodes # --------------------------- def initialize_db(state: CodeAssistantState): try: for _, row in extracted_data.iterrows(): embedding = embedding_model.encode([row['prompt'].strip()])[0] doc_id = hashlib.md5(row['prompt'].encode()).hexdigest() collection.add( documents=[row['canonical_solution'].strip()], metadatas=[{"prompt": row['prompt'], "type": "code_example"}], ids=[doc_id], embeddings=[embedding] ) return state except Exception as e: state["error"] = f"DB initialization failed: {str(e)}" return state def retrieve_examples(state: CodeAssistantState): try: embedding = embedding_model.encode([state["user_input"]])[0] results = collection.query( query_embeddings=[embedding], n_results=2 ) state["similar_examples"] = results['documents'][0] if results['documents'] else None return state except Exception as e: state["error"] = f"Retrieval failed: {str(e)}" return state def classify_task_llm(state: CodeAssistantState) -> CodeAssistantState: if not isinstance(state, dict): raise ValueError("State must be a dictionary") if "user_input" not in state or not state["user_input"].strip(): state["error"] = "No user input provided for classification" state["task_type"] = "generate" # Default to code generation return state try: prompt = f"""You are a helpful code assistant. Classify the user request as one of the following tasks: - "generate": if the user wants to write or generate code - "explain": if the user wants to understand what a code snippet does - "test": if the user wants to test existing code Return ONLY a JSON object in the format: {{"task": "...", "user_input": "..."}} — no explanation. User request: {state["user_input"]} """ completion = groq_client.chat.completions.create( model="llama3-70b-8192", messages=[ {"role": "system", "content": "Classify code-related user input. Respond with ONLY JSON."}, {"role": "user", "content": prompt} ], temperature=0.3, max_tokens=200, response_format={"type": "json_object"} ) content = completion.choices[0].message.content.strip() try: result = json.loads(content) if not isinstance(result, dict): raise ValueError("Response is not a JSON object") except (json.JSONDecodeError, ValueError) as e: state["error"] = f"Invalid response format from LLM: {str(e)}. Content: {content}" state["task_type"] = "generate" # Fallback to code generation return state task_type = result.get("task", "").lower() if task_type not in ["generate", "explain", "test"]: state["error"] = f"Invalid task type received: {task_type}" task_type = "generate" # Default to generation state["task_type"] = task_type state["user_input"] = result.get("user_input", state["user_input"]) return state except Exception as e: state["error"] = f"LLM-based classification failed: {str(e)}" state["task_type"] = "generate" # Fallback to code generation return state def test_code(state: CodeAssistantState) -> CodeAssistantState: if not isinstance(state, dict): raise ValueError("State must be a dictionary") if "user_input" not in state or not state["user_input"].strip(): state["error"] = "Please provide the code you want to test" return state try: messages = [ {"role": "system", "content": """You are a Python testing expert. Generate unit tests for the provided code. Return the test code in the following format: ```python # Test code here ```"""}, {"role": "user", "content": f"Generate comprehensive unit tests for this Python code:\n\n{state['user_input']}"} ] completion = groq_client.chat.completions.create( model="llama-3.3-70b-versatile", messages=messages, temperature=0.5, max_tokens=2048, ) test_code = completion.choices[0].message.content if test_code.startswith('```python'): test_code = test_code[9:-3] if test_code.endswith('```') else test_code[9:] elif test_code.startswith('```'): test_code = test_code[3:-3] if test_code.endswith('```') else test_code[3:] state["generated_tests"] = test_code.strip() state["metadata"] = { "model": "llama-3.3-70b-versatile", "timestamp": datetime.datetime.now().isoformat() } # Execute the tests and capture results try: # Create a temporary file to store the original code with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as code_file: code_file.write(state['user_input']) code_file_path = code_file.name # Create a temporary file to store the test code with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as test_file: test_file.write(test_code) test_file_path = test_file.name # Run the tests and capture output result = subprocess.run( ['python', test_file_path], capture_output=True, text=True, timeout=10 ) state["test_results"] = { "returncode": result.returncode, "stdout": result.stdout, "stderr": result.stderr } # Clean up temporary files os.unlink(code_file_path) os.unlink(test_file_path) except Exception as e: state["test_error"] = f"Error executing tests: {str(e)}" print(f"\nGenerated Tests:\n{test_code.strip()}\n") if "test_results" in state: print(f"Test Execution Results:\n{state['test_results']['stdout']}") if state["test_results"]["stderr"]: print(f"Errors:\n{state['test_results']['stderr']}") return state except Exception as e: state["error"] = f"Error generating tests: {str(e)}" return state def generate_code(state: CodeAssistantState) -> CodeAssistantState: if not isinstance(state, dict): raise ValueError("State must be a dictionary") if "user_input" not in state or not state["user_input"].strip(): state["error"] = "Please enter your code request" return state try: messages = [ {"role": "system", "content": "You are a Python coding assistant. Return only clean, production-ready code."}, {"role": "user", "content": state["user_input"].strip()} ] completion = groq_client.chat.completions.create( model="llama-3.3-70b-versatile", messages=messages, temperature=0.7, max_tokens=2048, ) code = completion.choices[0].message.content if code.startswith('```python'): code = code[9:-3] if code.endswith('```') else code[9:] elif code.startswith('```'): code = code[3:-3] if code.endswith('```') else code[3:] state["generated_code"] = code.strip() state["metadata"] = { "model": "llama-3.3-70b-versatile", "timestamp": datetime.datetime.now().isoformat() } # سطر طباعة النتيجة المضافة print(f"\nGenerated Code:\n{code.strip()}\n") return state except Exception as e: state["error"] = f"Error generating code: {str(e)}" return state def explain_code(state: CodeAssistantState) -> CodeAssistantState: try: messages = [ {"role": "system", "content": "You are a Python expert. Explain what the following code does in plain language."}, {"role": "user", "content": state["user_input"].strip()} ] completion = groq_client.chat.completions.create( model="llama-3.3-70b-versatile", messages=messages, temperature=0.5, max_tokens=1024 ) explanation = completion.choices[0].message.content.strip() state["generated_code"] = explanation state["metadata"] = { "model": "llama-3.3-70b-versatile", "timestamp": datetime.datetime.now().isoformat() } # سطر طباعة النتيجة المضافة print(f"Explanation:\n{explanation}") return state except Exception as e: state["error"] = f"Error explaining code: {str(e)}" return state # --------------------------- # 4. Build StateGraph Workflow (محدث) # --------------------------- workflow = StateGraph(CodeAssistantState) # إضافة جميع العقد بما فيها العقدة الجديدة workflow.add_node("initialize_db", initialize_db) workflow.add_node("retrieve_examples", retrieve_examples) workflow.add_node("classify_task", classify_task_llm) workflow.add_node("generate_code", generate_code) workflow.add_node("explain_code", explain_code) workflow.add_node("test_code", test_code) # العقدة الجديدة # تحديد نقطة البداية والروابط الأساسية workflow.set_entry_point("initialize_db") workflow.add_edge("initialize_db", "retrieve_examples") workflow.add_edge("retrieve_examples", "classify_task") # تحديث الروابط الشرطية لتشمل خيار الاختبار workflow.add_conditional_edges( "classify_task", lambda state: state["task_type"], { "generate": "generate_code", "explain": "explain_code", "test": "test_code" # الرابط الجديد } ) # إضافة روابط النهاية لجميع العقد workflow.add_edge("generate_code", END) workflow.add_edge("explain_code", END) workflow.add_edge("test_code", END) # الرابط الجديد # تجميع التدفق النهائي app_workflow = workflow.compile() # --------------------------- # 5. Create Gradio Interface # --------------------------- def process_input(user_input: str): """Function that will be called by Gradio to process user input""" initial_state = { "user_input": user_input, "similar_examples": None, "generated_code": None, "error": None, "task_type": None } result = app_workflow.invoke(initial_state) if result.get("error"): return f"Error: {result['error']}" if result["task_type"] == "generate": return f"Generated Code:\n\n{result['generated_code']}" else: return f"Code Explanation:\n\n{result['generated_code']}" # تعريف واجهة Gradio # Define Gradio interface with gr.Blocks(title="Smart Code Assistant") as demo: gr.Markdown(""" # Smart Code Assistant Enter your request either to generate new code or to explain existing code """) with gr.Row(): input_text = gr.Textbox(label="Enter your request", placeholder="Example: Write a function to add two numbers... or Explain this code...") output_text = gr.Textbox(label="Result", interactive=False) submit_btn = gr.Button("Execute") submit_btn.click(fn=process_input, inputs=input_text, outputs=output_text) # Quick examples gr.Examples( examples=[ ["Write a Python function to add two numbers"], ["Explain this code: for i in range(5): print(i)"], ["Create a function to convert temperature from Fahrenheit to Celsius"], ["test for i in range(3): print('Hello from test', i)"] ], inputs=input_text ) # تشغيل الواجهة if __name__ == "__main__": demo.launch()