import gradio as gr
import os
import spaces # Import the spaces library
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
import torch
from threading import Thread
import logging
from typing import Tuple, List, Dict, Generator
# --- Logging Configuration ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# --- Model & Quantization Settings ---
MODEL_ID = "unsloth/DeepSeek-R1-Distill-Qwen-7B-unsloth-bnb-4bit"
models: Dict[str, AutoModelForCausalLM] = {}
tokenizers: Dict[str, AutoTokenizer] = {}
bnb_config_4bit = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16, # Or torch.float16 if needed
)
def get_model_and_tokenizer() -> Tuple[AutoModelForCausalLM, AutoTokenizer]:
"""
Lazy-load the model and tokenizer if not already loaded.
Returns:
Tuple[model, tokenizer]: The loaded model and tokenizer.
"""
if "7B" not in models:
logging.info(f"Loading 7B model: {MODEL_ID} on demand")
try:
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
quantization_config=bnb_config_4bit,
torch_dtype=torch.bfloat16, # Or torch.float16 if needed
device_map='auto',
trust_remote_code=True,
)
model.eval() # Set the model to evaluation mode
models["7B"] = model
tokenizers["7B"] = tokenizer
logging.info("Loaded 7B model on demand.")
except Exception as e:
logging.error(f"Failed to load model and tokenizer: {e}")
raise e
return models["7B"], tokenizers["7B"]
# --- Default Prompt Templates for Multiple Presets ---
default_prompts = {
"coding": {
"brainstorm": """**Coding Brainstorm (Round 1)**
As a Senior Code Analyst, analyze the following problem and list key challenges and potential approaches.
**User Request:**
{user_prompt}
**Guidelines:**
1. Identify coding challenges.
2. Suggest potential methods and approaches.
3. Highlight any critical edge cases.
""",
"round2": """**Advanced Reasoning & Code Generation (Round 2)**
Based on your initial analysis:
**Initial Analysis:**
{brainstorm_response}
**User Request:**
{user_prompt}
**Task:**
1. Generate production-ready code with advanced reasoning.
2. Include a pun-filled birthday message with a coding twist within your output.
3. Comment the code clearly.
""",
"synthesis": """**Synthesis & Final Refinement (Round 3)**
Review the detailed code and reasoning below, and synthesize a final, refined response that:
1. Combines the brainstorming insights and advanced code generation.
2. Summarizes the solution succinctly.
3. Provides any additional improvements.
**Detailed Code & Reasoning:**
{round2_response}
""",
"rationale": """**Pun Generation and Rationale (Round 4)**
Based on the final refined response below, generate a clear, stand-alone pun-filled birthday message with a coding twist, then explain in detail why that pun was chosen.
Final Refined Response:
{final_response}
Your answer should:
1. Clearly output the pun as a separate line.
2. Explain the pun’s connection to birthdays and coding concepts (e.g., binary, syntax).
3. Describe any creative insights behind the choice.
"""
},
"math": {
"brainstorm": """**Math Problem Brainstorm (Round 1)**
As an expert mathematician, analyze the following problem and outline key concepts and strategies.
**Problem:**
{user_prompt}
**Guidelines:**
1. Identify the mathematical concepts involved.
2. List potential strategies or methods.
3. Note any assumptions or conditions.
""",
"round2": """**Solution Strategy Development (Round 2)**
Based on the initial analysis:
**Initial Analysis:**
{brainstorm_response}
**Problem:**
{user_prompt}
**Task:**
1. Develop a detailed strategy to solve the problem.
2. Include potential methods and intermediate steps.
""",
"synthesis": """**Solution Synthesis (Round 3)**
Review the strategy and previous analysis below, and produce a refined, step-by-step solution that:
1. Clearly explains the solution path.
2. Highlights key steps and justifications.
3. Summarizes the final answer.
**Detailed Strategy:**
{round2_response}
""",
"rationale": """**Solution Rationale (Round 4)**
Based on the final refined solution below, provide a detailed explanation of the key steps and mathematical insights.
Final Refined Solution:
{final_response}
Your response should:
1. Clearly explain why each step was taken.
2. Detail any assumptions and mathematical principles used.
3. Summarize the creative reasoning behind the solution.
"""
},
"writing": {
"brainstorm": """**Creative Brainstorm (Round 1)**
As a seasoned writer, brainstorm creative ideas for the following writing prompt.
**Writing Prompt:**
{user_prompt}
**Guidelines:**
1. List key themes and creative directions.
2. Suggest multiple approaches to the narrative.
3. Highlight any unique stylistic ideas.
""",
"round2": """**Outline Generation (Round 2)**
Based on the brainstorming below:
**Brainstormed Ideas:**
{brainstorm_response}
**Writing Prompt:**
{user_prompt}
**Task:**
1. Generate a detailed outline for a creative piece.
2. Organize the ideas into a coherent structure.
3. Provide bullet points or sections for the narrative.
""",
"synthesis": """**Draft Writing (Round 3)**
Review the outline below and produce a refined draft of the creative piece that:
1. Synthesizes the brainstorming insights and the outline.
2. Provides a coherent and engaging narrative.
3. Includes stylistic and thematic elements.
**Outline:**
{round2_response}
""",
"rationale": """**Final Editing and Rationale (Round 4)**
Based on the final draft below, refine the piece further and provide a detailed explanation of your creative choices.
Final Draft:
{final_response}
Your answer should:
1. Present the final refined text.
2. Explain the narrative choices, stylistic decisions, and thematic connections.
3. Detail any creative insights that influenced the final version.
"""
}
}
# --- Domain Detection ---
def detect_domain(user_prompt: str) -> str:
"""
Detect the domain based on keywords.
Args:
user_prompt (str): The user query.
Returns:
str: One of 'math', 'writing', or 'coding' (defaulting to coding).
"""
prompt_lower = user_prompt.lower()
math_keywords = ["solve", "integral", "derivative", "equation", "proof", "calculate", "sum", "product"]
writing_keywords = ["write", "story", "essay", "novel", "poem", "article", "narrative", "creative"]
coding_keywords = ["code", "program", "debug", "compile", "algorithm", "function"]
if any(kw in prompt_lower for kw in math_keywords):
logging.info("Domain detected as: math")
return "math"
elif any(kw in prompt_lower for kw in writing_keywords):
logging.info("Domain detected as: writing")
return "writing"
elif any(kw in prompt_lower for kw in coding_keywords):
logging.info("Domain detected as: coding")
return "coding"
else:
logging.info("No specific domain detected; defaulting to coding")
return "coding"
# --- Memory Management ---
class MemoryManager:
"""Encapsulate shared memory for storing and retrieving conversation items."""
def __init__(self) -> None:
self.shared_memory: List[str] = []
def store(self, item: str) -> None:
"""Store a memory item and log an excerpt."""
self.shared_memory.append(item)
logging.info(f"[Memory Stored]: {item[:50]}...")
def retrieve(self, query: str, top_k: int = 3) -> List[str]:
"""Retrieve recent memory items containing the query text."""
query_lower = query.lower()
relevant = [item for item in self.shared_memory if query_lower in item.lower()]
if not relevant:
logging.info("[Memory Retrieval]: No relevant memories found.")
else:
logging.info(f"[Memory Retrieval]: Found {len(relevant)} relevant memories.")
return relevant[-top_k:]
global_memory_manager = MemoryManager()
# --- Unified Generation Function ---
def generate_response(model, tokenizer, prompt: str, max_tokens: int, temperature: float, top_p: float) -> str:
"""Generate a response for a given prompt."""
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
kwargs = dict(
input_ids=input_ids,
streamer=streamer,
max_new_tokens=max_tokens,
temperature=temperature,
top_p=top_p,
do_sample=True,
)
thread = Thread(target=model.generate, kwargs=kwargs)
with torch.no_grad():
thread.start()
response = ""
try:
for text in streamer:
response += text
except Exception as e:
logging.error(f"Error during generation: {e}")
raise e
thread.join()
return response
# --- Multi-Round Agent Class ---
class MultiRoundAgent:
"""
Encapsulate the multi-round prompt chaining and response generation.
This class runs a 4-round pipeline based on the given preset.
"""
def __init__(self, model, tokenizer, prompt_templates: Dict[str, str], memory_manager: MemoryManager):
self.model = model
self.tokenizer = tokenizer
self.prompt_templates = prompt_templates
self.memory_manager = memory_manager
def run_pipeline(self, user_prompt: str, params: Dict, show_raw: bool = False) -> Generator[str, None, None]:
# Round 1: Brainstorming / Analysis
logging.info("--- Round 1 ---")
prompt_r1 = self.prompt_templates["brainstorm"].format(user_prompt=user_prompt)
r1 = generate_response(self.model, self.tokenizer, prompt_r1, params.get("max_new_tokens"), params.get("temp"), params.get("top_p"))
self.memory_manager.store(f"Round 1 Response: {r1}")
# Round 2: Secondary Generation (strategy/outline/code)
logging.info("--- Round 2 ---")
prompt_r2 = self.prompt_templates["round2"].format(brainstorm_response=r1, user_prompt=user_prompt)
r2 = generate_response(self.model, self.tokenizer, prompt_r2, params.get("max_new_tokens") + 100, params.get("temp"), params.get("top_p"))
self.memory_manager.store(f"Round 2 Response: {r2}")
# Round 3: Synthesis & Refinement (streaming updates)
logging.info("--- Round 3 ---")
prompt_r3 = self.prompt_templates["synthesis"].format(round2_response=r2)
input_ids_r3 = self.tokenizer.encode(prompt_r3, return_tensors="pt").to(self.model.device)
streamer_r3 = TextIteratorStreamer(self.tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
kwargs_r3 = dict(
input_ids=input_ids_r3,
streamer=streamer_r3,
max_new_tokens=params.get("max_new_tokens") // 2,
temperature=params.get("temp"),
top_p=params.get("top_p")
)
thread_r3 = Thread(target=self.model.generate, kwargs=kwargs_r3)
with torch.no_grad():
thread_r3.start()
r3 = ""
try:
for text in streamer_r3:
r3 += text
yield r3 # Yield progressive updates from Round 3
except Exception as e:
logging.error(f"Error during Round 3 streaming: {e}")
raise e
thread_r3.join()
self.memory_manager.store(f"Final Synthesis Response: {r3}")
# Round 4: Rationale / Final Output
logging.info("--- Round 4 ---")
prompt_r4 = self.prompt_templates["rationale"].format(final_response=r3)
r4 = generate_response(self.model, self.tokenizer, prompt_r4, 300, params.get("temp"), params.get("top_p"))
self.memory_manager.store(f"Round 4 Response: {r4}")
# Construct final output based on the show_raw flag.
if show_raw:
final_output = (
f"{r4}\n\n[Raw Outputs]\n"
f"Round 1:\n{r1}\n\n"
f"Round 2:\n{r2}\n\n"
f"Round 3:\n{r3}\n\n"
f"Round 4:\n{r4}\n"
)
else:
final_output = r4
yield final_output
# --- Swarm Agent Iterative Function ---
@spaces.GPU(duration=180) # Adjust duration as needed
def swarm_agent_iterative(user_prompt: str, temp: float, top_p: float, max_new_tokens: int, memory_top_k: int,
prompt_templates: Dict[str, str], domain: str, show_raw: bool) -> Generator[str, None, None]:
"""
Wraps the multi-round agent functionality. Depending on the detected domain,
it runs the 4-round pipeline.
"""
model, tokenizer = get_model_and_tokenizer()
agent = MultiRoundAgent(model, tokenizer, prompt_templates, global_memory_manager)
params = {"temp": temp, "top_p": top_p, "max_new_tokens": max_new_tokens}
return agent.run_pipeline(user_prompt, params, show_raw)
# --- Explanation Function for Additional Requests ---
def handle_explanation_request(user_prompt: str, history: List) -> str:
"""
Retrieve stored rationale and additional context from conversation history,
then generate an explanation.
"""
retrieved = global_memory_manager.retrieve("Round 4 Response:", top_k=3)
explanation_prompt = "Below are previous final outputs and related context from our conversation:\n"
if retrieved:
for item in retrieved:
explanation_prompt += f"- {item}\n"
else:
explanation_prompt += "No stored final output found.\n"
explanation_prompt += "\nRecent related exchanges:\n"
for chat in history:
if ("explain" in chat[0].lower()) or (chat[1] and "explain" in chat[1].lower()):
explanation_prompt += f"User: {chat[0]}\nAssistant: {chat[1]}\n"
explanation_prompt += "\nBased on the above context, please provide a detailed explanation of the creative choices."
model, tokenizer = get_model_and_tokenizer()
explanation = generate_response(model, tokenizer, explanation_prompt, 300, 0.7, 0.9)
return explanation
# --- Helper to Format History ---
def format_history(history: List) -> List[Dict[str, str]]:
"""
Convert history (list of [user, assistant] pairs) into a list of message dictionaries.
"""
messages = []
for item in history:
if isinstance(item, (list, tuple)) and len(item) == 2:
user_msg, assistant_msg = item
messages.append({"role": "user", "content": user_msg})
if assistant_msg:
messages.append({"role": "assistant", "content": assistant_msg})
elif isinstance(item, dict):
messages.append(item)
return messages
# --- Gradio Chat Interface Function ---
def gradio_interface(message: str, history: List, param_state: Dict, prompt_state: Dict) -> Generator[List[Dict[str, str]], None, None]:
"""
Called by Gradio's ChatInterface. Uses current generation parameters and preset prompt templates.
If the user asks for an explanation, routes the request accordingly.
"""
if "explain" in message.lower():
explanation = handle_explanation_request(message, history)
history = history + [[message, explanation]]
yield format_history(history)
return
try:
temp = float(param_state.get("temperature", 0.5))
top_p = float(param_state.get("top_p", 0.9))
max_new_tokens = int(param_state.get("max_new_tokens", 300))
memory_top_k = int(param_state.get("memory_top_k", 2))
show_raw = bool(param_state.get("show_raw_output", False))
except Exception as e:
logging.error(f"Parameter conversion error: {e}")
temp, top_p, max_new_tokens, memory_top_k, show_raw = 0.5, 0.9, 300, 2, False
domain = detect_domain(message)
# Get the prompt templates for the detected domain; default to coding if not set.
prompt_templates = prompt_state.get(domain, default_prompts.get(domain, default_prompts["coding"]))
history = history + [[message, ""]]
for partial_response in swarm_agent_iterative(
user_prompt=message,
temp=temp,
top_p=top_p,
max_new_tokens=max_new_tokens,
memory_top_k=memory_top_k,
prompt_templates=prompt_templates,
domain=domain,
show_raw=show_raw
):
history[-1][1] = partial_response
yield format_history(history)
# --- UI Settings & Styling ---
ui_description = '''
DeepSeek Agent Swarm Chat
Multi-round agent with 4-round prompt chaining for three presets:
- Coding
- Math
- Writing
'''
ui_license = """
---
"""
ui_placeholder = """
DeepSeek Agent Swarm
Ask me anything...
"""
css = """
h1 {
text-align: center;
display: block;
}
#duplicate-button {
margin: auto;
color: white;
background: #1565c0;
border-radius: 100vh;
}
"""
# --- Gradio UI ---
with gr.Blocks(css=css, title="DeepSeek Agent Swarm Chat") as demo:
gr.Markdown(ui_description)
# Hidden states for parameters and prompt configurations.
param_state = gr.State({
"temperature": 0.5,
"top_p": 0.9,
"max_new_tokens": 300,
"memory_top_k": 2,
"show_raw_output": False, # New parameter for raw output
})
prompt_state = gr.State({
"coding": default_prompts["coding"],
"math": default_prompts["math"],
"writing": default_prompts["writing"],
})
with gr.Tabs():
with gr.Tab("Chat"):
chatbot = gr.Chatbot(height=450, placeholder=ui_placeholder, label="Agent Swarm Output", type="messages")
gr.ChatInterface(
fn=gradio_interface,
chatbot=chatbot,
additional_inputs=[param_state, prompt_state],
examples=[
['How can we build a robust web service that scales efficiently under load?'],
['Solve the integral of x^2 from 0 to 1.'],
['Write a short story about a mysterious writer in a busy city.'],
['Create a pun-filled birthday message with a coding twist.']
],
cache_examples=False,
type="messages",
)
with gr.Tab("Parameters"):
gr.Markdown("### Generation Parameters")
temp_slider = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.5, label="Temperature")
top_p_slider = gr.Slider(minimum=0.01, maximum=1.0, step=0.05, value=0.9, label="Top P")
max_tokens_num = gr.Number(value=300, label="Max new tokens", precision=0)
memory_topk_slider = gr.Slider(minimum=1, maximum=5, step=1, value=2, label="Memory Retrieval Top K")
show_raw_checkbox = gr.Checkbox(value=False, label="Show Raw Output") # New checkbox for raw output
save_params_btn = gr.Button("Save Parameters")
save_params_btn.click(
lambda t, p, m, k, s: {
"temperature": t,
"top_p": p,
"max_new_tokens": m,
"memory_top_k": k,
"show_raw_output": s
},
inputs=[temp_slider, top_p_slider, max_tokens_num, memory_topk_slider, show_raw_checkbox],
outputs=param_state,
)
with gr.Tab("Prompt Config"):
gr.Markdown("### Configure Prompt Templates for Each Preset")
with gr.Tabs():
with gr.Tab("Coding"):
prompt_brainstorm_box_code = gr.Textbox(
value=default_prompts["coding"]["brainstorm"],
label="Brainstorm Prompt (Coding)",
lines=8,
)
prompt_round2_box_code = gr.Textbox(
value=default_prompts["coding"]["round2"],
label="Round 2 Prompt (Coding)",
lines=8,
)
prompt_synthesis_box_code = gr.Textbox(
value=default_prompts["coding"]["synthesis"],
label="Synthesis Prompt (Coding)",
lines=8,
)
prompt_rationale_box_code = gr.Textbox(
value=default_prompts["coding"]["rationale"],
label="Rationale Prompt (Coding)",
lines=8,
)
with gr.Tab("Math"):
prompt_brainstorm_box_math = gr.Textbox(
value=default_prompts["math"]["brainstorm"],
label="Brainstorm Prompt (Math)",
lines=8,
)
prompt_round2_box_math = gr.Textbox(
value=default_prompts["math"]["round2"],
label="Round 2 Prompt (Math)",
lines=8,
)
prompt_synthesis_box_math = gr.Textbox(
value=default_prompts["math"]["synthesis"],
label="Synthesis Prompt (Math)",
lines=8,
)
prompt_rationale_box_math = gr.Textbox(
value=default_prompts["math"]["rationale"],
label="Rationale Prompt (Math)",
lines=8,
)
with gr.Tab("Writing"):
prompt_brainstorm_box_writing = gr.Textbox(
value=default_prompts["writing"]["brainstorm"],
label="Brainstorm Prompt (Writing)",
lines=8,
)
prompt_round2_box_writing = gr.Textbox(
value=default_prompts["writing"]["round2"],
label="Round 2 Prompt (Writing)",
lines=8,
)
prompt_synthesis_box_writing = gr.Textbox(
value=default_prompts["writing"]["synthesis"],
label="Synthesis Prompt (Writing)",
lines=8,
)
prompt_rationale_box_writing = gr.Textbox(
value=default_prompts["writing"]["rationale"],
label="Rationale Prompt (Writing)",
lines=8,
)
save_prompts_btn = gr.Button("Save Prompts")
def save_prompts(code_brain, code_r2, code_syn, code_rat, math_brain, math_r2, math_syn, math_rat, writing_brain, writing_r2, writing_syn, writing_rat):
return {
"coding": {
"brainstorm": code_brain,
"round2": code_r2,
"synthesis": code_syn,
"rationale": code_rat,
},
"math": {
"brainstorm": math_brain,
"round2": math_r2,
"synthesis": math_syn,
"rationale": math_rat,
},
"writing": {
"brainstorm": writing_brain,
"round2": writing_r2,
"synthesis": writing_syn,
"rationale": writing_rat,
}
}
save_prompts_btn.click(
save_prompts,
inputs=[prompt_brainstorm_box_code, prompt_round2_box_code, prompt_synthesis_box_code, prompt_rationale_box_code,
prompt_brainstorm_box_math, prompt_round2_box_math, prompt_synthesis_box_math, prompt_rationale_box_math,
prompt_brainstorm_box_writing, prompt_round2_box_writing, prompt_synthesis_box_writing, prompt_rationale_box_writing],
outputs=prompt_state,
)
gr.Markdown(ui_license)
if __name__ == "__main__":
demo.launch()