{
    "cells": [
        {
            "cell_type": "markdown",
            "id": "9f97dd1e",
            "metadata": {},
            "source": [
                "# Libraries"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 145,
            "id": "9e974df6",
            "metadata": {},
            "outputs": [],
            "source": [
                "import os\n",
                "import re\n",
                "from typing import TypedDict, List, Optional, Annotated\n",
                "from IPython.display import Image, display\n",
                "\n",
                "from langchain_core.documents import Document\n",
                "from langchain_core.messages import BaseMessage, SystemMessage\n",
                "from langchain_core.tools import tool\n",
                "from langgraph.checkpoint.memory import InMemorySaver\n",
                "from langgraph.graph.message import add_messages\n",
                "from langchain_ollama import ChatOllama, OllamaEmbeddings\n",
                "from langchain_elasticsearch import ElasticsearchStore\n",
                "from langgraph.graph import StateGraph, END\n",
                "from langgraph.prebuilt import ToolNode\n",
                "from langfuse import get_client, Langfuse\n",
                "from langfuse.langchain import CallbackHandler\n",
                "\n",
                "from typing import TypedDict, List, Optional, Annotated, Literal\n",
                "from pydantic import BaseModel, Field\n",
                "from langchain_core.messages import BaseMessage, SystemMessage, AIMessage"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "30edcecc",
            "metadata": {},
            "outputs": [],
            "source": [
                "ES_URL = os.getenv(\"ELASTICSEARCH_LOCAL_URL\")\n",
                "INDEX_NAME = os.getenv(\"ELASTICSEARCH_INDEX\")\n",
                "CODE_INDE\n",
                "BASE_URL = os.getenv(\"OLLAMA_LOCAL_URL\")\n",
                "MODEL_NAME = os.getenv(\"OLLAMA_MODEL_NAME\")\n",
                "EMB_MODEL_NAME = os.getenv(\"OLLAMA_EMB_MODEL_NAME\")\n",
                "LANGFUSE_PUBLIC_KEY = os.getenv(\"LANGFUSE_PUBLIC_KEY\")\n",
                "LANGFUSE_SECRET_KEY = os.getenv(\"LANGFUSE_SECRET_KEY\")\n",
                "LANGFUSE_HOST = os.getenv(\"LANGFUSE_HOST\")\n",
                "\n",
                "\n",
                "embeddings = OllamaEmbeddings(base_url=BASE_URL, model=EMB_MODEL_NAME)\n",
                "llm = ChatOllama(base_url=BASE_URL, model=MODEL_NAME, temperature=0)\n",
                "\n",
                "vector_store = ElasticsearchStore(\n",
                "    es_url=ES_URL,\n",
                "    index_name=INDEX_NAME,\n",
                "    embedding=embeddings,\n",
                "    query_field=\"text\",\n",
                "    vector_query_field=\"vector\",\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "873ea2f6",
            "metadata": {},
            "source": [
                "### State"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 147,
            "id": "5f8c88cf",
            "metadata": {},
            "outputs": [],
            "source": [
                "class AgentState(TypedDict):\n",
                "    messages: Annotated[list, add_messages]"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 148,
            "id": "30473bce",
            "metadata": {},
            "outputs": [],
            "source": [
                "class AgentResponse(BaseModel):\n",
                "    \"\"\"\n",
                "    Structured output contract for final assistant responses.\n",
                "    \"\"\"\n",
                "    language: Literal[\"en\"] = Field(\n",
                "        description=\"ISO code. Must always be 'en'.\"\n",
                "    )\n",
                "    content: str = Field(\n",
                "        description=\"Final answer in English.\"\n",
                "    )"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "1d60c120",
            "metadata": {},
            "source": [
                "### Tools"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 149,
            "id": "f0a21230",
            "metadata": {},
            "outputs": [],
            "source": [
                "retrieve_kwargs = {\"k\": 5}"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 150,
            "id": "f9359747",
            "metadata": {},
            "outputs": [],
            "source": [
                "def format_context(docs: List[Document]) -> str:\n",
                "    chunks: List[str] = []\n",
                "    for i, doc in enumerate(docs, 1):\n",
                "        source = (doc.metadata or {}).get(\"source\", \"Untitled\")\n",
                "        source_id = (doc.metadata or {}).get(\"id\", f\"chunk-{i}\")\n",
                "        text = doc.page_content or \"\"\n",
                "        chunks.append(f\"[{i}] id={source_id} source={source}\\n{text}\")\n",
                "    return \"\\n\\n\".join(chunks)\n",
                "\n",
                "\n",
                "@tool\n",
                "def retrieve(query: str) -> str:\n",
                "    \"\"\"This tool retrieves relevant documents from the vector store based on the input query and formats them for the agent's response.\n",
                "    Args:\n",
                "        query (str): The input query for which to retrieve relevant documents.\n",
                "    \"\"\"\n",
                "    retriever = vector_store.as_retriever(\n",
                "        search_type=\"similarity\",\n",
                "        search_kwargs=retrieve_kwargs,\n",
                "    )\n",
                "    docs = retriever.invoke(query)\n",
                "    return format_context(docs)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 151,
            "id": "e5247ab9",
            "metadata": {},
            "outputs": [],
            "source": [
                "def should_continue(state: AgentState) -> str:\n",
                "    last = state[\"messages\"][-1]\n",
                "    \n",
                "    if getattr(last, \"tool_calls\", None):\n",
                "        return \"tools\"\n",
                "    return \"end\""
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 152,
            "id": "a644f6fa",
            "metadata": {},
            "outputs": [],
            "source": [
                "tools = [retrieve]\n",
                "tool_node = ToolNode(tools)\n",
                "memory = InMemorySaver()"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "395966e2",
            "metadata": {},
            "source": [
                "### Agent"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 153,
            "id": "36d0f54e",
            "metadata": {},
            "outputs": [],
            "source": [
                "def _message_text(message: BaseMessage) -> str:\n",
                "    content = getattr(message, \"content\", \"\")\n",
                "    if isinstance(content, str):\n",
                "        return content\n",
                "    if isinstance(content, list):\n",
                "        return \" \".join(str(item) for item in content)\n",
                "    return str(content)\n",
                "\n",
                "\n",
                "def _last_user_query(messages: List[BaseMessage]) -> str:\n",
                "    for message in reversed(messages):\n",
                "        message_type = getattr(message, \"type\", \"\")\n",
                "        if message_type == \"human\":\n",
                "            return _message_text(message)\n",
                "    return _message_text(messages[-1]) if messages else \"\"\n",
                "\n",
                "\n",
                "def _last_assistant_text(messages: List[BaseMessage]) -> str:\n",
                "    for message in reversed(messages):\n",
                "        message_type = getattr(message, \"type\", \"\")\n",
                "        if message_type == \"ai\":\n",
                "            return _message_text(message)\n",
                "    return \"\"\n",
                "\n",
                "\n",
                "MAX_LANGUAGE_RETRIES = 5\n",
                "\n",
                "\n",
                "def agent(state: AgentState) -> AgentState:\n",
                "    messages: List[BaseMessage] = state[\"messages\"]\n",
                "    user_query = _last_user_query(messages)\n",
                "\n",
                "    retrieved_context = retrieve.invoke({\"query\": user_query})\n",
                "\n",
                "    system = SystemMessage(\n",
                "        content=(\n",
                "            \"\"\"\n",
                "            You are an AVAP language assistant focused on accurate and grounded answers.\n",
                "\n",
                "            Rules:\n",
                "            1. ALWAYS use the provided retrieval context from Elasticsearch FIRST.\n",
                "            2. Use ONLY the retrieved context to produce the final answer. Do NOT invent AVAP commands.\n",
                "            3. ALWAYS reply in English only. Never respond in any other language.\n",
                "            4. If asked for code/snippet, include exactly ONE fenced AVAP block (```avap ... ```).\n",
                "            5. Never output internal reasoning or unrelated text.\n",
                "\n",
                "            Retrieved context:\n",
                "            \"\"\"\n",
                "            + retrieved_context\n",
                "        )\n",
                "    )\n",
                "\n",
                "\n",
                "\n",
                "    structured_model = llm.with_structured_output(AgentResponse)\n",
                "    structured_response: AgentResponse = structured_model.invoke(\n",
                "        [system, *messages]\n",
                "    )\n",
                "\n",
                "    # Keep LangGraph message flow unchanged\n",
                "    final_message = AIMessage(content=structured_response.content)\n",
                "    return {\"messages\": [*messages, final_message]}\n",
                "\n",
                "    # model = llm.bind_tools(tools)\n",
                "    # response = model.invoke([system, *messages])\n",
                "\n",
                "    # return {\"messages\": [*messages, response]}"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "ef55bca3",
            "metadata": {},
            "source": [
                "### Graph"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 154,
            "id": "fae46a58",
            "metadata": {},
            "outputs": [],
            "source": [
                "graph = StateGraph(AgentState)\n",
                "graph.add_node(\"agent\", agent)\n",
                "\n",
                "\n",
                "graph.set_entry_point(\"agent\")\n",
                "graph.add_edge(\"agent\", END)\n",
                "\n",
                "# Alternative mode (single pass) - kept commented for quick rollback.\n",
                "# graph.set_entry_point(\"agent\")\n",
                "# graph.add_edge(\"agent\", END)\n",
                "\n",
                "agent_graph = graph.compile()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 155,
            "id": "2fec3fdb",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "image/png": "[encoded data, 8296 characters]",
                        "text/plain": [
                            "<IPython.core.display.Image object>"
                        ]
                    },
                    "metadata": {},
                    "output_type": "display_data"
                }
            ],
            "source": [
                "try:\n",
                "    display(Image(agent_graph.get_graph().draw_mermaid_png()))\n",
                "except Exception:\n",
                "    pass"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "1e9aff05",
            "metadata": {},
            "source": [
                "### Test"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 156,
            "id": "8569cf39",
            "metadata": {},
            "outputs": [],
            "source": [
                "config = {\"configurable\": {\"thread_id\": \"5\"}, \n",
                "          #\"callbacks\": [langfuse_handler],\n",
                "          #\"run_name\": \"rag-local-test\"differences between getDatetime() and getTimeStamp() functions in AVAP\n",
                "          }\n",
                "\n",
                "def stream_graph_updates(user_input: str):\n",
                "    for event in agent_graph.stream(\n",
                "        {\"messages\": [{\"role\": \"user\", \"content\": user_input}]},\n",
                "        #config=config,\n",
                "        stream_mode=\"values\",\n",
                "    ):\n",
                "        event[\"messages\"][-1].pretty_print()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 162,
            "id": "a1a1f3cf",
            "metadata": {},
            "outputs": [],
            "source": [
                "user_input = (\n",
                "    \"Generate two variables, asigning them one number to each, do it in AVAP language.\"\n",
                ")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 163,
            "id": "53b89690",
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "================================\u001b[1m Human Message \u001b[0m=================================\n",
                        "\n",
                        "Generate two variables, asigning them one number to each, do it in AVAP language.\n",
                        "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
                        "\n",
                        "user: generate two variables, assigning them one number to each, you assign themize in, avap language.\n"
                    ]
                }
            ],
            "source": [
                "a = stream_graph_updates(user_input)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "136f420c",
            "metadata": {},
            "outputs": [],
            "source": []
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "assistance-engine",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.12.11"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}
