app.py

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')

# ---------------------------
# 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()