src/rag_pipelines.py

from __future__ import annotations

import json
from typing import List, Tuple

from openai import OpenAI

from src.prompts import SYSTEM_PROMPT
from src.ticketing import create_github_issue


def retrieve_documents(query, db, k=3):
    return db.similarity_search(query, k=k)


def build_context(docs) -> Tuple[str, List[str]]:
    context_parts: List[str] = []
    sources: List[str] = []

    for doc in docs:
        context_parts.append(doc.page_content)
        source = doc.metadata.get("source", "unknown")
        page = doc.metadata.get("page", "?")
        source_info = f"{source} - page {page}"
        if source_info not in sources:
            sources.append(source_info)

    return "\n\n".join(context_parts), sources


def format_chat_history(memory):
    messages = []
    for item in memory:
        messages.append({"role": "user", "content": item["user"]})
        messages.append({"role": "assistant", "content": item["bot"]})
    return messages


def _ticket_tool():
    return [
        {
            "type": "function",
            "function": {
                "name": "create_support_ticket",
                "description": "Create a support ticket when user has an unresolved issue.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "title": {"type": "string"},
                        "description": {"type": "string"},
                    },
                    "required": ["title", "description"],
                },
            },
        }
    ]


def generate_answer(query, context, sources, memory):
    client = OpenAI()
    history_messages = format_chat_history(memory)

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            *history_messages,
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}"},
        ],
        tools=_ticket_tool(),
        tool_choice="auto",
    )

    message = response.choices[0].message

    if message.tool_calls:
        tool_call = message.tool_calls[0]
        if tool_call.function.name == "create_support_ticket":
            args = json.loads(tool_call.function.arguments)
            issue_url = create_github_issue(
                title=args["title"],
                description=args["description"],
            )
            return f"Support ticket result: {issue_url}"

    answer = message.content or "I could not generate an answer."
    if not sources:
        return answer
    return answer + "\n\nSources:\n" + "\n".join(sources)