main.py

import json
import time
from typing import List, Dict, Any

from flask import Flask, request, jsonify, Response
from flask_cors import CORS
from dotenv import load_dotenv
from ollama import chat

from assistant_config import (
    MODEL_NAME,
    OLLAMA_OPTIONS,
    ENABLE_WEB_SEARCH,
    MAX_SEARCH_RESULTS,
    SYSTEM_PROMPT,
    WEB_SEARCH_DECISION_PROMPT,
    ANSWER_WITH_SEARCH_PROMPT,
    ANSWER_WITHOUT_SEARCH_PROMPT,
)

# Load environment variables from .env
load_dotenv()


# =============================
# Ollama Web Search Import
# =============================

try:
    from ollama import web_search
    WEB_SEARCH_AVAILABLE = True
except Exception:
    WEB_SEARCH_AVAILABLE = False


# =============================
# Flask App
# =============================

app = Flask(__name__)
CORS(app)


# =============================
# Helper Functions
# =============================

def get_last_user_message(messages: List[Dict[str, Any]]) -> str:
    """
    Get the latest user message from the message list sent by AnythingLLM.
    """
    for message in reversed(messages):
        if message.get("role") == "user":
            return message.get("content", "")
    return ""


def should_use_web_search(question: str) -> bool:
    """
    Ask the LLM to decide whether the user's question requires Web Search.
    The decision prompt is stored in assistant_config.py.
    """

    if not ENABLE_WEB_SEARCH:
        return False

    if not WEB_SEARCH_AVAILABLE:
        print("[Web Search Decision] Web Search tool is not available.")
        return False

    judge_prompt = WEB_SEARCH_DECISION_PROMPT.format(question=question)

    try:
        response = chat(
            model=MODEL_NAME,
            messages=[
                {
                    "role": "user",
                    "content": judge_prompt
                }
            ],
            options={
                "temperature": 0,
                "top_p": 0.1,
                "num_ctx": 2048,
            }
        )

        decision = response.message.content.strip().upper()

        print(f"[Web Search Decision Raw] {decision}")

        return decision.startswith("YES")

    except Exception as e:
        print(f"[Web Search Decision Error] {e}")
        return False


def run_web_search(question: str) -> str:
    """
    Run Ollama Web Search and convert the search results into a text context.
    """

    if not WEB_SEARCH_AVAILABLE:
        return ""

    try:
        results = web_search(
            query=question,
            max_results=MAX_SEARCH_RESULTS
        )

        search_context = ""

        for i, item in enumerate(results.results, start=1):
            title = getattr(item, "title", "No title")
            url = getattr(item, "url", "No URL")
            content = getattr(item, "content", "")

            search_context += f"[{i}] {title}\n"
            search_context += f"URL: {url}\n"
            search_context += f"Summary: {content}\n\n"

        return search_context

    except Exception as e:
        print(f"[Web Search Error] {e}")
        return f"Web Search failed: {e}"


def build_messages(request_messages: List[Dict[str, Any]], user_question: str, search_context: str = "") -> List[Dict[str, str]]:
    """
    Build the final message list that will be sent to the Ollama model.
    Preserve prior user/assistant history and inject the configured system prompt.
    """

    # Keep conversation history except for the final user query, which we may modify for search.
    last_user_index = None
    for idx in range(len(request_messages) - 1, -1, -1):
        if request_messages[idx].get("role") == "user":
            last_user_index = idx
            break

    conversation_history: List[Dict[str, str]] = []
    if last_user_index is not None:
        for message in request_messages[:last_user_index]:
            role = message.get("role")
            if role in {"user", "assistant"}:
                conversation_history.append({
                    "role": role,
                    "content": message.get("content", "")
                })

    if search_context:
        final_user_content = ANSWER_WITH_SEARCH_PROMPT.format(
            search_context=search_context,
            user_question=user_question
        )
    else:
        final_user_content = ANSWER_WITHOUT_SEARCH_PROMPT.format(
            user_question=user_question
        )

    return [
        {
            "role": "system",
            "content": SYSTEM_PROMPT
        },
        *conversation_history,
        {
            "role": "user",
            "content": final_user_content
        }
    ]


def ask_ollama(request_messages: List[Dict[str, Any]]) -> str:
    """
    Main response workflow:
    1. Ask the LLM to decide whether Web Search is needed.
    2. If needed, run Ollama Web Search.
    3. Send the search results and the user's question and history to the Ollama model.
    """

    user_question = get_last_user_message(request_messages)
    search_context = ""

    needs_search = should_use_web_search(user_question)

    print(f"[Web Search Decision] {needs_search}")

    if needs_search:
        print("[Web Search] Searching the web...")
        search_context = run_web_search(user_question)

        if search_context:
            print("[Web Search] Search completed.")
        else:
            print("[Web Search] No search context returned.")
    else:
        print("[Web Search] Skipped.")

    messages = build_messages(
        request_messages=request_messages,
        user_question=user_question,
        search_context=search_context
    )

    response = chat(
        model=MODEL_NAME,
        messages=messages,
        options=OLLAMA_OPTIONS
    )

    return response.message.content


def openai_stream_response(answer: str, model_name: str):
    """
    Return an OpenAI-compatible Server-Sent Events response.
    This prevents AnythingLLM from waiting forever when it sends stream=true.
    """

    created_time = int(time.time())
    completion_id = f"chatcmpl-{created_time}"

    # First chunk: assistant role
    role_chunk = {
        "id": completion_id,
        "object": "chat.completion.chunk",
        "created": created_time,
        "model": model_name,
        "choices": [
            {
                "index": 0,
                "delta": {
                    "role": "assistant"
                },
                "finish_reason": None
            }
        ]
    }

    yield f"data: {json.dumps(role_chunk, ensure_ascii=False)}\n\n"

    # Content chunk
    content_chunk = {
        "id": completion_id,
        "object": "chat.completion.chunk",
        "created": created_time,
        "model": model_name,
        "choices": [
            {
                "index": 0,
                "delta": {
                    "content": answer
                },
                "finish_reason": None
            }
        ]
    }

    yield f"data: {json.dumps(content_chunk, ensure_ascii=False)}\n\n"

    # Final chunk
    done_chunk = {
        "id": completion_id,
        "object": "chat.completion.chunk",
        "created": created_time,
        "model": model_name,
        "choices": [
            {
                "index": 0,
                "delta": {},
                "finish_reason": "stop"
            }
        ]
    }

    yield f"data: {json.dumps(done_chunk, ensure_ascii=False)}\n\n"
    yield "data: [DONE]\n\n"


# =============================
# OpenAI-compatible Endpoints
# =============================

@app.route("/", methods=["GET"])
def root():
    """
    Health check endpoint.
    """
    return jsonify({
        "message": "Assistant API is running.",
        "model": MODEL_NAME,
        "web_search_enabled": ENABLE_WEB_SEARCH,
        "web_search_available": WEB_SEARCH_AVAILABLE
    })


@app.route("/v1/models", methods=["GET"])
def list_models():
    """
    Return a model list in an OpenAI-compatible format.
    AnythingLLM can use this endpoint to detect available models.
    """
    return jsonify({
        "object": "list",
        "data": [
            {
                "id": MODEL_NAME,
                "object": "model",
                "created": int(time.time()),
                "owned_by": "local-ollama"
            }
        ]
    })


@app.route("/v1/chat/completions", methods=["POST", "GET"])
def chat_completions():
    """
    OpenAI-compatible chat completions endpoint.
    Supports both normal JSON responses and streaming responses.
    """
    # If a browser or user visits this URL with GET, return a helpful message
    if request.method == "GET":
        return jsonify({
            "message": "This endpoint accepts POST with an OpenAI-compatible `messages` payload.",
            "usage": {
                "endpoint": "/v1/chat/completions",
                "method": "POST",
                "example": {
                    "messages": [{"role": "user", "content": "What should I eat in Taipei?"}],
                    "stream": False
                }
            }
        }), 200

    data = request.get_json(force=True)

    messages = data.get("messages", [])
    stream = data.get("stream", False)

    print(f"[Request] stream={stream}")

    user_question = get_last_user_message(messages)

    if not user_question:
        answer = "Please enter a question."
    else:
        answer = ask_ollama(messages)

    created_time = int(time.time())
    completion_id = f"chatcmpl-{created_time}"

    # If AnythingLLM requests streaming, return SSE format.
    if stream:
        return Response(
            openai_stream_response(answer, MODEL_NAME),
            mimetype="text/event-stream"
        )

    # Normal non-streaming JSON response.
    return jsonify({
        "id": completion_id,
        "object": "chat.completion",
        "created": created_time,
        "model": MODEL_NAME,
        "choices": [
            {
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": answer
                },
                "finish_reason": "stop"
            }
        ],
        "usage": {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0
        }
    })


# =============================
# Run Flask App
# =============================

if __name__ == "__main__":
    app.run(
        host="127.0.0.1",
        port=8000,
        debug=True,
        threaded=True
    )