main.py

import gradio as gr
import os
import json
from openai import OpenAI

# Настройки API
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "dummy")
API_BASE_URL = os.getenv("API_BASE_URL", "http://127.0.0.1:8088/v1")

# Инициализация клиента OpenAI
client = OpenAI(
    api_key=OPENAI_API_KEY,
    base_url=API_BASE_URL
)

def fetch_models():
    """Fetches available models from the API"""
    try:
        models = client.models.list()
        model_ids = [model.id for model in models.data]
        if not model_ids:
            raise ValueError("No models available from API")
        return model_ids
    except Exception as e:
        raise RuntimeError(f"Failed to fetch models from API: {str(e)}") from e

def parse_streaming_json(buffer):
    """Extracts complete JSON objects from streaming buffer
    
    Args:
        buffer: String containing partial or complete JSON objects
        
    Returns:
        tuple: (list of parsed JSON objects, remaining buffer)
    """
    parsed_objects = []
    remaining = buffer
    
    while remaining.strip():
        # Find potential JSON object boundaries
        brace_count = 0
        start_idx = None
        
        for i, char in enumerate(remaining):
            if char == '{':
                if brace_count == 0:
                    start_idx = i
                brace_count += 1
            elif char == '}':
                brace_count -= 1
                if brace_count == 0 and start_idx is not None:
                    # Found complete JSON object
                    json_str = remaining[start_idx:i+1]
                    try:
                        obj = json.loads(json_str)
                        parsed_objects.append(obj)
                        remaining = remaining[i+1:]
                        break
                    except json.JSONDecodeError:
                        # Not valid JSON yet, continue
                        continue
        else:
            # No complete JSON found, return what we have
            break
    
    return parsed_objects, remaining

def format_tool_output(tool_data):
    """Formats tool output as readable Markdown
    
    Args:
        tool_data: Parsed JSON object from tool
        
    Returns:
        str: Formatted Markdown string
    """
    tool_type = tool_data.get("tool_name_discriminator", "")
    
    if tool_type == "generateplantool":
        # Format research plan
        goal = tool_data.get("research_goal", "N/A")
        steps = tool_data.get("planned_steps", [])
        strategies = tool_data.get("search_strategies", [])
        
        result = "## План исследования\n\n"
        result += f"**Цель:** {goal}\n\n"
        
        if steps:
            result += "### Запланированные шаги:\n\n"
            for i, step in enumerate(steps, 1):
                result += f"{i}. {step}\n"
            result += "\n"
        
        if strategies:
            result += "### Стратегии поиска:\n\n"
            for strategy in strategies:
                result += f"- {strategy}\n"
            result += "\n"
        
        return result
    
    elif tool_type == "adaptplantool":
        # Format plan adaptation
        new_goal = tool_data.get("new_goal", "")
        changes = tool_data.get("plan_changes", [])
        next_steps = tool_data.get("next_steps", [])
        
        result = "## Адаптация плана\n\n"
        
        if new_goal:
            result += f"**Обновленная цель:** {new_goal}\n\n"
        
        if changes:
            result += "### Изменения в плане:\n\n"
            for change in changes:
                result += f"- {change}\n"
            result += "\n"
        
        if next_steps:
            result += "### Следующие шаги:\n\n"
            for i, step in enumerate(next_steps, 1):
                result += f"{i}. {step}\n"
            result += "\n"
        
        return result
    
    elif tool_type == "finalanswertool":
        # Return only the answer content
        answer = tool_data.get("answer", "")
        if answer:
            return f"## Результат исследования\n\n{answer}\n"
        return "## Результат исследования\n\nОтвет не найден.\n"
    
    else:
        # For unknown tool types or intermediate data, show abbreviated info
        # Check if this looks like search results
        if "Search Query" in str(tool_data) or "Search Results" in str(tool_data):
            result = "## Поиск информации\n\n"
            result += "*Выполняется поиск по источникам...*\n\n"
            return result
        
        # Default: show tool type
        return f"*Обработка: {tool_type or 'unknown tool'}...*\n\n"

def format_tool_display(tool_data):
    """Formats tool output with structured display
    
    Args:
        tool_data: Parsed JSON object from tool
        
    Returns:
        str: Formatted output with goal as header, steps in expanded block, and JSON in collapsed block
    """
    tool_name = tool_data.get("tool_name_discriminator", "unknown")
    result = ""
    
    # DEBUG: show what we received
    import sys
    print(f"DEBUG: tool_name={tool_name}, keys={list(tool_data.keys())}", file=sys.stderr)
    
    # Special handling for finalanswertool
    if tool_name == "finalanswertool":
        # 1. Collapsed reasoning block
        reasoning = tool_data.get("reasoning", "")
        if reasoning:
            result += f"<details>\n<summary>Reasoning</summary>\n\n{reasoning}\n\n</details>\n\n"
        
        # 2. Expanded final answer block
        answer = tool_data.get("answer", "")
        if answer:
            result += f"<details open>\n<summary>Final answer</summary>\n\n{answer}\n\n</details>\n\n"
        
        # No JSON block for finalanswertool - answer already displayed above
        return result
    
    # Standard handling for other tools
    # 1. research_goal as header
    if "research_goal" in tool_data:
        result += f"## Цель: {tool_data['research_goal']}\n\n"
    
    # 2. Steps in expanded block
    completed = tool_data.get("completed_steps", [])
    planned = tool_data.get("planned_steps", [])
    
    if completed or planned:
        result += "<details open>\n<summary>Прогресс выполнения</summary>\n\n"
        
        if completed:
            result += "**Выполненные шаги:**\n"
            for step in completed:
                result += f"- {step}\n"
            result += "\n"
        
        if planned:
            result += "**Запланированные шаги:**\n"
            for step in planned:
                result += f"- {step}\n"
        
        result += "</details>\n\n"
    
    # 3. Collapsed block with full JSON
    json_str = json.dumps(tool_data, indent=2, ensure_ascii=False)
    result += f"<details>\n<summary>Tool: {tool_name}</summary>\n\n```json\n{json_str}\n```\n\n</details>\n\n"
    
    return result

def create_collapsible_block(title, content, content_type="json"):
    """Creates a collapsible HTML details block with formatted content
    
    Args:
        title: Title for the summary line
        content: Content to display (will be formatted as JSON or text)
        content_type: Type of content - 'json' or 'text'
        
    Returns:
        str: HTML details block
    """
    if content_type == "json":
        formatted_content = f"```json\n{content}\n```"
    else:
        formatted_content = content
    
    return f'''<details>
<summary>{title}</summary>

{formatted_content}

</details>'''

def extract_answer_from_json(text):
    """Extracts answer field from JSON response if present"""
    try:
        # Try to parse as JSON
        data = json.loads(text)
        
        # If finalanswertool with answer field, return formatted answer
        if data.get("tool_name_discriminator") == "finalanswertool" and "answer" in data:
            return data["answer"]
        
        # Otherwise return pretty-printed JSON for debugging
        return f"```json\n{json.dumps(data, ensure_ascii=False, indent=2)}\n```"
    
    except json.JSONDecodeError:
        # If not valid JSON, return as is
        return text

def chat_completion(user_message, model, chat_history=None, agent_id=None):
    """Sends request to OpenAI API and returns response with streaming
    
    Args:
        user_message: User's message text
        model: Model ID to use (or agent_id for continuing session)
        chat_history: List of previous messages in OpenAI format
        agent_id: Current agent ID (if continuing session)
        
    Returns:
        Generator yielding tuples: (output_text, updated_history, agent_id, clarification_questions)
    """
    import sys
    
    try:
        # Initialize or update chat history
        if chat_history is None:
            chat_history = []
        
        # Add user message to history
        chat_history.append({"role": "user", "content": user_message})
        
        # Use agent_id if provided, otherwise use selected model
        model_to_use = agent_id if agent_id else model
        
        stream = client.chat.completions.create(
            model=model_to_use,
            messages=chat_history,
            stream=True,
            temperature=0.2
        )
        
        # Buffers for accumulating tool call arguments and output blocks
        tool_calls_buffer = {}
        output_blocks = []
        processed_tool_ids = set()
        block_counter = 0
        has_final_answer = False
        current_agent_id = agent_id
        clarification_questions = []
        content_buffer = ""  # Buffer for accumulating delta.content
        
        for chunk in stream:
            if not chunk.choices:
                continue
            
            # Extract agent ID from model field
            if chunk.model and chunk.model.startswith("sgr_agent_"):
                current_agent_id = chunk.model
            
            delta = chunk.choices[0].delta
            
            # Process tool calls
            if delta.tool_calls:
                for tool_call in delta.tool_calls:
                    # Get tool call index (used to track partial updates)
                    tool_idx = tool_call.index if hasattr(tool_call, 'index') else 0
                    
                    # Initialize buffer for this tool call if needed
                    if tool_idx not in tool_calls_buffer:
                        tool_calls_buffer[tool_idx] = {
                            'id': tool_call.id if hasattr(tool_call, 'id') else None,
                            'name': '',
                            'arguments': ''
                        }
                    
                    # Accumulate function name and arguments
                    if tool_call.function:
                        if tool_call.function.name:
                            tool_calls_buffer[tool_idx]['name'] = tool_call.function.name
                        if tool_call.function.arguments:
                            tool_calls_buffer[tool_idx]['arguments'] += tool_call.function.arguments
                    
                    # Try to parse if we have complete JSON
                    args_str = tool_calls_buffer[tool_idx]['arguments']
                    tool_id = tool_calls_buffer[tool_idx].get('id', tool_idx)
                    tool_name = tool_calls_buffer[tool_idx].get('name', '')
                    
                    if args_str and tool_id not in processed_tool_ids:
                        try:
                            # Parse JSON to check if complete
                            tool_data = json.loads(args_str)
                            
                            # Check if this is clarification request
                            if tool_name == "clarification":
                                clarification_questions = tool_data.get("questions", [])
                                processed_tool_ids.add(tool_id)
                                # Don't add to output blocks, will be handled separately
                                continue
                            
                            # Check if this is finalanswertool
                            if tool_data.get("tool_name_discriminator") == "finalanswertool":
                                has_final_answer = True
                                print(f"DEBUG: Found finalanswertool, has_final_answer={has_final_answer}", file=sys.stderr)
                            
                            # Format tool output with structured display
                            block_counter += 1
                            block = format_tool_display(tool_data)
                            print(f"DEBUG: Formatted block length={len(block)}, first 100 chars={block[:100]}", file=sys.stderr)
                            output_blocks.append(block)
                            processed_tool_ids.add(tool_id)
                            
                            # Yield all accumulated blocks
                            yield ("\n\n".join(output_blocks), chat_history, current_agent_id, clarification_questions)
                        except json.JSONDecodeError:
                            # Arguments not complete yet, continue accumulating
                            pass
            
            # Process regular content (skip if we already have final answer to avoid duplication)
            if delta.content and not has_final_answer:
                print(f"DEBUG: Got content chunk, length={len(delta.content)}", file=sys.stderr)
                
                # Accumulate content in buffer
                content_buffer += delta.content
                
                # Try to parse accumulated content as JSON
                if '{' in content_buffer and '"tool_name_discriminator"' in content_buffer:
                    print(f"DEBUG: Content buffer looks like JSON, attempting to parse (buffer length={len(content_buffer)})", file=sys.stderr)
                    
                    # Try to find complete JSON objects
                    parsed_objects, remaining = parse_streaming_json(content_buffer)
                    
                    for tool_data in parsed_objects:
                        if tool_data.get("tool_name_discriminator") == "finalanswertool":
                            print(f"DEBUG: Found finalanswertool in content buffer!", file=sys.stderr)
                            has_final_answer = True
                            block_counter += 1
                            block = format_tool_display(tool_data)
                            output_blocks.append(block)
                            content_buffer = remaining  # Keep unparsed remainder
                            yield ("\n\n".join(output_blocks), chat_history, current_agent_id, clarification_questions)
                            break
                    
                    if has_final_answer:
                        continue
                
                # If not JSON or not parsed yet, create regular content block
                # But only for complete chunks (not partial JSON)
                if not ('{' in content_buffer and content_buffer.count('{') > content_buffer.count('}')):
                    block_counter += 1
                    block = create_collapsible_block(
                        title=f"Content (#{block_counter})",
                        content=content_buffer,
                        content_type="text"
                    )
                    output_blocks.append(block)
                    content_buffer = ""  # Reset buffer
                    yield ("\n\n".join(output_blocks), chat_history, current_agent_id, clarification_questions)
        
        # Final yield to ensure all content is displayed
        if output_blocks or clarification_questions:
            yield ("\n\n".join(output_blocks), chat_history, current_agent_id, clarification_questions)
    
    except Exception as e:
        yield (f"Request error: {str(e)}", chat_history or [], None, [])

def send_clarification(clarification_text, agent_id, chat_history):
    """Sends clarification response to continue research session
    
    Args:
        clarification_text: User's clarification/answer text
        agent_id: Agent ID from previous request
        chat_history: Current chat history
        
    Returns:
        Generator yielding tuples: (output_text, updated_history, agent_id, clarification_questions)
    """
    if not agent_id:
        yield ("Error: No active agent session", chat_history, None, [])
        return
    
    # Use chat_completion with agent_id to continue session
    yield from chat_completion(clarification_text, None, chat_history, agent_id)

def handle_submit(user_message, model, chat_history, agent_id):
    """Wrapper for chat_completion that manages UI state
    
    Returns:
        Generator yielding: (output, history, agent_id, questions_md, group_visible)
    """
    for output, history, new_agent_id, questions in chat_completion(user_message, model, chat_history, agent_id):
        # Format questions for display
        if questions:
            questions_md = "**Questions:**\n\n" + "\n".join(f"{i+1}. {q}" for i, q in enumerate(questions))
            yield output, history, new_agent_id, questions_md, gr.update(visible=True)
        else:
            yield output, history, new_agent_id, "", gr.update(visible=False)

def handle_clarification(clarification_text, agent_id, chat_history):
    """Wrapper for send_clarification that manages UI state
    
    Returns:
        Generator yielding: (output, history, agent_id, questions_md, group_visible)
    """
    for output, history, new_agent_id, questions in send_clarification(clarification_text, agent_id, chat_history):
        # Format questions for display
        if questions:
            questions_md = "**Questions:**\n\n" + "\n".join(f"{i+1}. {q}" for i, q in enumerate(questions))
            yield output, history, new_agent_id, questions_md, gr.update(visible=True)
        else:
            yield output, history, new_agent_id, "", gr.update(visible=False)

# Fetch available models on startup
available_models = fetch_models()

# Create Gradio interface
with gr.Blocks(title="SGR Deep Research via OpenAI API") as interface:
    gr.Markdown("# SGR Deep Research via OpenAI API")
    gr.Markdown("Simple interface for working with ChatGPT")
    
    # State variables
    chat_history_state = gr.State([])
    agent_id_state = gr.State(None)
    
    model_dropdown = gr.Dropdown(
        choices=available_models,
        value=available_models[0],
        label="Model"
    )
    
    user_input = gr.Textbox(
        label="Your request",
        placeholder="Enter your message...",
        lines=3
    )
    
    output = gr.Markdown(
        label="Response"
    )
    
    submit_btn = gr.Button("Submit")
    
    # Clarification section
    with gr.Group(visible=False) as clarification_group:
        gr.Markdown("## Clarification needed")
        clarification_questions_display = gr.Markdown("")
        clarification_input = gr.Textbox(
            label="Your clarification",
            placeholder="Provide additional details...",
            lines=2
        )
        clarification_btn = gr.Button("Send clarification")
    
    # Wire up submit button
    submit_btn.click(
        fn=handle_submit,
        inputs=[user_input, model_dropdown, chat_history_state, agent_id_state],
        outputs=[output, chat_history_state, agent_id_state, clarification_questions_display, clarification_group]
    )
    
    # Wire up clarification button
    clarification_btn.click(
        fn=handle_clarification,
        inputs=[clarification_input, agent_id_state, chat_history_state],
        outputs=[output, chat_history_state, agent_id_state, clarification_questions_display, clarification_group]
    )

if __name__ == "__main__":
    interface.launch()