Update app.py

app.py

@@ -1,154 +1,182 @@
 import gradio as gr
-import numpy as np
-import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
-import torch
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
-
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
-
-    return image, seed
-
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
-
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-
-            run_button = gr.Button("Run", scale=0, variant="primary")
-
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
+import requests
+from pydantic import BaseModel
+import time
+import json
+import os
+from typing import Generator, Tuple, List
+
+class StepResponse(BaseModel):
+    title: str
+    content: str
+    next_action: str
+    confidence: float
+
+def get_available_models() -> List[str]:
+    """Fetch available models from OpenRouter API"""
+    headers = {
+        "Authorization": f"Bearer {os.getenv('OPENROUTER_API_KEY')}",
+    }
+    
+    try:
+        response = requests.get("https://openrouter.ai/api/v1/models", headers=headers)
+        response.raise_for_status()
+        models = response.json()
+        return [model["id"] for model in models]
+    except Exception as e:
+        print(f"Error fetching models: {e}")
+        # Fallback to a basic list of known models
+        return [
+            "anthropic/claude-3-sonnet-20240320",
+            "anthropic/claude-3-opus-20240229",
+            "google/gemini-pro",
+            "meta-llama/llama-2-70b-chat",
+            "mistral/mistral-medium",
+        ]
+
+def make_api_call(model: str, system_prompt: str, messages: list, max_tokens: int, 
+                 is_final_answer: bool = False) -> StepResponse:
+    """Make API call to OpenRouter with specified model"""
+    headers = {
+        "HTTP-Referer": "https://localhost:7860",  # Gradio default
+        "X-Title": "Reasoning Chain Demo",
+        "Authorization": f"Bearer {os.getenv('OPENROUTER_API_KEY')}",
+        "Content-Type": "application/json"
+    }
+    
+    url = "https://openrouter.ai/api/v1/chat/completions"
+    
+    request_body = {
+        "model": model,
+        "messages": [
+            {"role": "system", "content": system_prompt},
+            *messages
+        ],
+        "max_tokens": max_tokens,
+        "temperature": 0.2,
+        "response_format": {"type": "json_object"}
+    }
+    
+    for attempt in range(3):
+        try:
+            response = requests.post(url, headers=headers, json=request_body)
+            response.raise_for_status()
+            
+            result = response.json()
+            message_content = result['choices'][0]['message']['content']
+            
+            try:
+                response_data = json.loads(message_content)
+                return StepResponse(**response_data)
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Failed to parse JSON response: {str(e)}")
+                
+        except Exception as e:
+            if attempt == 2:
+                return StepResponse(
+                    title="Error",
+                    content=f"Failed to generate {'final answer' if is_final_answer else 'step'} after 3 attempts. Error: {str(e)}",
+                    next_action="final_answer",
+                    confidence=0.5
                 )
-
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
+            time.sleep(1)
+
+def generate_response(prompt: str, model: str, progress=gr.Progress()) -> Generator[str, None, None]:
+    """Generator function that yields formatted markdown for each step"""
+    system_prompt = """You are an AI assistant that explains your reasoning step by step, incorporating dynamic Chain of Thought (CoT), reflection, and verbal reinforcement learning. IMPORTANT: You must output exactly ONE step of reasoning at a time:
+
+1. Each response must contain ONE single step of your reasoning process.
+2. For each step, enclose your thoughts within <thinking> tags as you explore that specific step.
+3. After completing your current step, indicate whether you need another step or are ready for the final answer.
+4. Do not try to complete multiple steps or the entire analysis in one response.
+5. Regularly evaluate your progress, being critical and honest about your reasoning process.
+6. Assign a quality score between 0.0 and 1.0 to guide your approach:
+   - 0.8+: Continue current approach
+   - 0.5-0.7: Consider minor adjustments
+   - Below 0.5: Seriously consider backtracking and trying a different approach
+
+IMPORTANT: Your response must be a valid JSON object with the following structure:
+{
+    "title": "Step title or topic",
+    "content": "Detailed step content",
+    "next_action": "One of: continue, reflect, or final_answer",
+    "confidence": float between 0.0 and 1.0
+}"""
+
+    messages = [{"role": "user", "content": prompt}]
+    step_count = 1
+    markdown_output = ""
+    
+    while True:
+        progress(step_count / 15, f"Step {step_count}")  # Show progress
+        step_data = make_api_call(model, system_prompt, messages, 750)
+        
+        # Format step as markdown
+        step_md = f"### Step {step_count}: {step_data.title}\n\n"
+        step_md += f"{step_data.content}\n\n"
+        step_md += f"**Confidence:** {step_data.confidence:.2f}\n\n"
+        step_md += "---\n\n"
+        
+        markdown_output += step_md
+        yield markdown_output  # Update the output incrementally
+        
+        messages.append({"role": "assistant", "content": json.dumps(step_data.model_dump())})
+        
+        if step_data.next_action == 'final_answer' and step_count < 15:
+            messages.append({"role": "user", "content": "Please continue your analysis with at least 5 more steps before providing the final answer."})
+        elif step_data.next_action == 'final_answer':
+            break
+        elif step_data.next_action == 'reflect' or step_count % 3 == 0:
+            messages.append({"role": "user", "content": "Please perform a detailed self-reflection on your reasoning so far."})
+        else:
+            messages.append({"role": "user", "content": "Please continue with the next step in your analysis."})
+        
+        step_count += 1
+
+    # Generate final answer
+    final_data = make_api_call(model, system_prompt, messages, 750, is_final_answer=True)
+    
+    final_md = f"### Final Answer\n\n"
+    final_md += f"{final_data.content}\n\n"
+    final_md += f"**Confidence:** {final_data.confidence:.2f}\n\n"
+    
+    markdown_output += final_md
+    yield markdown_output
+
+def create_interface():
+    # Check for API key
+    if not os.getenv('OPENROUTER_API_KEY'):
+        raise ValueError("Please set OPENROUTER_API_KEY environment variable")
+    
+    available_models = get_available_models()
+    
+    with gr.Blocks() as interface:
+        gr.Markdown("# AI Reasoning Chain with Model Selection")
+        gr.Markdown("This demo shows chain-of-thought reasoning across different language models.")
+        
+        with gr.Row():
+            with gr.Column():
+                model_dropdown = gr.Dropdown(
+                    choices=available_models,
+                    value=available_models[0],
+                    label="Select Model"
                 )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2,  # Replace with defaults that work for your model
+                query_input = gr.Textbox(
+                    label="Enter your query:",
+                    placeholder="e.g., What are the potential long-term effects of climate change on global agriculture?"
                 )
-
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, seed],
-    )
+                submit_btn = gr.Button("Generate Response")
+            
+            output_box = gr.Markdown(label="Response")
+        
+        submit_btn.click(
+            fn=generate_response,
+            inputs=[query_input, model_dropdown],
+            outputs=output_box
+        )
+    
+    return interface
 
 if __name__ == "__main__":
-    demo.launch()
+    interface = create_interface()
+    interface.launch()