{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9f97dd1e",
   "metadata": {},
   "source": [
    "# Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9e974df6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from typing import TypedDict, List, Optional, Annotated, Literal\n",
    "from IPython.display import Image, display\n",
    "\n",
    "from langchain_core.documents import Document\n",
    "from langchain_core.messages import BaseMessage, SystemMessage, AIMessage, ToolMessage\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_elasticsearch import ElasticsearchStore\n",
    "from langgraph.graph import StateGraph, END\n",
    "from langgraph.prebuilt import ToolNode, tools_condition\n",
    "from langfuse import Langfuse\n",
    "from langfuse.decorators import observe, langfuse_context\n",
    "\n",
    "from src.utils.llm_factory import create_chat_model\n",
    "from src.utils.emb_factory import create_embedding_model\n",
    "from src.config import settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "30edcecc",
   "metadata": {},
   "outputs": [],
   "source": [
    "langfuse = Langfuse()\n",
    "\n",
    "llm = create_chat_model(\n",
    "    provider=\"ollama\",\n",
    "    model=settings.ollama_model_name,\n",
    "    temperature=0.5,\n",
    "    validate_model_on_init=True,\n",
    ")\n",
    "embeddings = create_embedding_model(\n",
    "    provider=\"ollama\",\n",
    "    model=settings.ollama_emb_model_name,\n",
    ")\n",
    "vector_store = ElasticsearchStore(\n",
    "    es_url=settings.elasticsearch_local_url,\n",
    "    index_name=\"avap-docs-test-v3\",\n",
    "    embedding=embeddings,\n",
    "    query_field=\"text\",\n",
    "    vector_query_field=\"embedding\",\n",
    "    # strategy=ElasticsearchStore.ApproxRetrievalStrategy(\n",
    "    #     hybrid=True,\n",
    "    #     rrf={\"rank_constant\": 60, \"window_size\": 100}\n",
    "    # )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ad98841b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Langfuse client is authenticated and ready!\n"
     ]
    }
   ],
   "source": [
    "if langfuse.auth_check():\n",
    "    print(\"Langfuse client is authenticated and ready!\")\n",
    "else:\n",
    "    print(\"Authentication failed. Please check your credentials and host.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "873ea2f6",
   "metadata": {},
   "source": [
    "### State"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5f8c88cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AgentState(TypedDict):\n",
    "    messages: Annotated[list, add_messages]\n",
    "    reformulated_query: str\n",
    "    context: str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fd8ed542",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AgenticAgentState(TypedDict):\n",
    "    messages: Annotated[list, add_messages]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d60c120",
   "metadata": {},
   "source": [
    "### Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f0a21230",
   "metadata": {},
   "outputs": [],
   "source": [
    "retrieve_kwargs = {\"k\": 3}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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",
    "@observe(name=\"context_retrieve\")\n",
    "def context_retrieve(query: str) -> str:\n",
    "    \"\"\"Consults vector store to respond AVAP related questions\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",
    "    context = format_context(docs)\n",
    "\n",
    "    langfuse_context.update_current_observation(\n",
    "        input={\"query\": query, \"k\": retrieve_kwargs[\"k\"]},\n",
    "        output={\n",
    "            \"documents_count\": len(docs),\n",
    "            \"sources\": [(doc.metadata or {}).get(\"source\", \"Untitled\") for doc in docs],\n",
    "            \"document_ids\": [(doc.metadata or {}).get(\"id\", f\"chunk-{i+1}\") for i, doc in enumerate(docs)],\n",
    "            \"context_preview\": context[:1000],\n",
    "        },\n",
    "    )\n",
    "    return context"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "395966e2",
   "metadata": {},
   "source": [
    "### Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "66ae23f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "REFORMULATE_PROMPT = SystemMessage(\n",
    "    content=(\n",
    "        \"You are a deterministic query rewriting function.\\n\"\n",
    "        \"You convert natural language questions into keyword search queries.\\n\\n\"\n",
    "        \"Strict constraints:\\n\"\n",
    "        \"1. Keep function names and technical tokens unchanged.\\n\"\n",
    "        \"2. Remove filler phrases.\\n\"\n",
    "        \"3. Do not answer.\\n\"\n",
    "        \"4. Do not explain.\\n\"\n",
    "        \"5. Do not generate code.\\n\"\n",
    "        \"6. Return a single-line query only.\\n\"\n",
    "        \"7. If already optimal, return unchanged.\\n\"\n",
    "    )\n",
    ")\n",
    "\n",
    "GENERATE_PROMPT = SystemMessage(\n",
    "    content=\"\"\"You are an agent designed to assist users with AVAP (Advanced Virtual API Programming) language.\n",
    "    It's a new language, so you should know nothing about it.\n",
    "    Use ONLY the provided context to answer AVAP-related questions.\n",
    "    If the context does not contain enough information, say so honestly.\n",
    "    If the question is not related to AVAP, answer based on your general knowledge.\n",
    "\n",
    "    Context:\n",
    "    {context}\"\"\"\n",
    ")\n",
    "\n",
    "AGENTIC_PROMPT = SystemMessage(\n",
    "    content=\"\"\"\n",
    "    You are an assistant that helps users with AVAP (Advanced Virtual API Programming) language questions.\n",
    "\n",
    "    AVAP is a completely new programming language and you have NO built-in knowledge about it.\n",
    "\n",
    "    Rules:\n",
    "\n",
    "    1. If the user question is related to AVAP:\n",
    "    - You must use the `context_retrieve` tool before answering.\n",
    "    - The tool output is INTERNAL CONTEXT, not a user message.\n",
    "    - Never thank the user for the retrieved context.\n",
    "    - You must synthesize an answer to respond the user query based SOLELY on the retrieved context.\n",
    "    \n",
    "    2. If the retrieved context is insufficient:\n",
    "    - Call `context_retrieve` again with a better reformulated query.\n",
    "\n",
    "    3. If the question is not related to AVAP:\n",
    "    - Answer normally using general knowledge.\n",
    "\n",
    "    4. If the user asks for code only:\n",
    "    - Return only the code snippet.\n",
    "    - Do not add explanation, headings, or markdown fences unless requested.\n",
    "    \"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "36d0f54e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def reformulate(state: AgentState) -> AgentState:\n",
    "    \"\"\"Use the LLM to rewrite the user query for better retrieval.\"\"\"\n",
    "    user_msg = state[\"messages\"][-1]\n",
    "    resp = llm.invoke([REFORMULATE_PROMPT, user_msg])\n",
    "    reformulated = resp.content.strip()\n",
    "    print(f\"[reformulate] '{user_msg.content}' → '{reformulated}'\")\n",
    "    return {\"reformulated_query\": reformulated}\n",
    "\n",
    "\n",
    "def retrieve(state: AgentState) -> AgentState:\n",
    "    \"\"\"Retrieve context using the reformulated query.\"\"\"\n",
    "    query = state[\"reformulated_query\"]\n",
    "    docs = vector_store.as_retriever(\n",
    "        search_type=\"similarity\",\n",
    "        search_kwargs=retrieve_kwargs,\n",
    "    ).invoke(query)\n",
    "    context = format_context(docs)\n",
    "    print(f\"[retrieve] {len(docs)} docs fetched\")\n",
    "    print(context)\n",
    "    return {\"context\": context}\n",
    "\n",
    "\n",
    "def generate(state: AgentState) -> AgentState:\n",
    "    \"\"\"Generate the final answer using retrieved context.\"\"\"\n",
    "    prompt = SystemMessage(\n",
    "        content=GENERATE_PROMPT.content.format(context=state[\"context\"])\n",
    "    )\n",
    "    resp = llm.invoke([prompt] + state[\"messages\"])\n",
    "    return {\"messages\": [resp]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f073edc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [context_retrieve]\n",
    "\n",
    "def agent(state: AgenticAgentState) -> AgenticAgentState:\n",
    "    llm_with_tools = llm.bind_tools(tools)\n",
    "    return {\"messages\": [llm_with_tools.invoke([SystemMessage(content=AGENTIC_PROMPT.content)] + state[\"messages\"])]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef55bca3",
   "metadata": {},
   "source": [
    "### Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "fae46a58",
   "metadata": {},
   "outputs": [],
   "source": [
    "memory = InMemorySaver()\n",
    "\n",
    "graph_builder = StateGraph(AgentState)\n",
    "\n",
    "graph_builder.add_node(\"reformulate\", reformulate)\n",
    "graph_builder.add_node(\"retrieve\", retrieve)\n",
    "graph_builder.add_node(\"generate\", generate)\n",
    "\n",
    "graph_builder.set_entry_point(\"reformulate\")\n",
    "graph_builder.add_edge(\"reformulate\", \"retrieve\")\n",
    "graph_builder.add_edge(\"retrieve\", \"generate\")\n",
    "graph_builder.add_edge(\"generate\", END)\n",
    "\n",
    "guided_graph = graph_builder.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7f57b543",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 14976 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "try:\n",
    "    display(Image(guided_graph.get_graph().draw_mermaid_png()))\n",
    "except Exception:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f7a0993f",
   "metadata": {},
   "outputs": [],
   "source": [
    "tool_node = ToolNode(tools=tools)\n",
    "memory = InMemorySaver()\n",
    "\n",
    "graph_builder = StateGraph(AgenticAgentState)\n",
    "\n",
    "graph_builder.add_node(\"agent\", agent)\n",
    "graph_builder.add_node(\"tools\", tool_node)\n",
    "\n",
    "graph_builder.set_entry_point(\"agent\")\n",
    "graph_builder.add_conditional_edges(\n",
    "    \"agent\",\n",
    "    tools_condition,\n",
    ")\n",
    "graph_builder.add_edge(\"tools\", \"agent\")\n",
    "\n",
    "agentic_graph = graph_builder.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2fec3fdb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 12472 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "try:\n",
    "    display(Image(agentic_graph.get_graph().draw_mermaid_png()))\n",
    "except Exception:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e9aff05",
   "metadata": {},
   "source": [
    "### Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "8569cf39",
   "metadata": {},
   "outputs": [],
   "source": [
    "@observe(name=\"graph_run\")\n",
    "def stream_graph_updates(user_input: str, graph: StateGraph):\n",
    "    langfuse_context.update_current_trace(\n",
    "        user_id=\"alberto\",\n",
    "        tags=[\"avap\", \"rag\", \"langgraph\"],\n",
    "        metadata={\"feature\": \"agentic-rag\"},\n",
    "    )\n",
    "\n",
    "    for event in graph.stream(\n",
    "        {\"messages\": [{\"role\": \"user\", \"content\": user_input}]},\n",
    "        stream_mode=\"values\",\n",
    "    ):\n",
    "        event[\"messages\"][-1].pretty_print()\n",
    "\n",
    "    return event[\"messages\"][-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a1a1f3cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "user_input = \"\"\"What types of includes does AVAP have?\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "53b89690",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "What types of includes does AVAP have?\n",
      "[reformulate] 'What types of includes does AVAP have?' → '\"avap includes type\"'\n",
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "What types of includes does AVAP have?\n",
      "[retrieve] 3 docs fetched\n",
      "[1] id=chunk-1 source=Untitled\n",
      "\n",
      "\n",
      "Token:\n",
      "\n",
      "\n",
      "\n",
      "ASSIGN\n",
      "\n",
      "[2] id=chunk-2 source=Untitled\n",
      "\n",
      "\n",
      "> **Nota de implementación:** `<connector_instantiation>` se distingue de `<orm_connector_init>` (ORM) únicamente por contexto semántico: el UUID pasado como argumento determina si el adaptador resuelto es un ORM de base de datos o un proxy de terceros. La gramática los trata de forma idéntica; el motor de ejecución selecciona el adaptador apropiado en runtime.\n",
      "\n",
      "---\n",
      "\n",
      "## SECCIÓN VI: Utilidades, Criptografía y Manipulación de Datos\n",
      "\n",
      "AVAP incluye un set de comandos integrados de alto nivel para manipular tipos complejos (JSON y Listas), tiempos, textos y generar hashes.\n",
      "\n",
      "### 6.1 Manipulación Nativa de Listas y Objetos JSON\n",
      "Para extraer y mutar estructuras complejas, AVAP provee comandos nativos específicos:\n",
      "* **`variableToList(elemento, destino)`**: Fuerza a que una variable escalar se convierta en una estructura iterable de lista.\n",
      "* **`itemFromList(lista_origen, indice, destino)`**: Extrae de forma segura el elemento contenido en la posición `indice` de una lista.\n",
      "* **`variableFromJSON(json_origen, clave, destino)`**: Parsea un objeto JSON en memoria y extrae el valor correspondiente a la `clave`.\n",
      "* **`AddVariableToJSON(clave, valor, json_destino)`**: Inyecta dinámicamente una nueva propiedad dentro de un objeto JSON existente.\n",
      "\n",
      "### 6.2 Criptografía y Expresiones Regulares\n",
      "* **`encodeSHA256` y `encodeMD5(origen, destino)`**: Funciones criptográficas que encriptan de forma irreversible un texto. Vitales para el almacenamiento seguro de contraseñas.\n",
      "* **`getRegex(origen, patron, destino)`**: Aplica una Expresión Regular (`patron`) sobre la variable de origen, extrayendo las coincidencias exactas.\n",
      "\n",
      "### 6.3 Transformación de Tiempo y Cadenas\n",
      "* **Fechas:** `getTimeStamp` (convierte un string a Epoch), `getDateTime` (Epoch a string legible), y `stampToDatetime` (Epoch a objeto datetime estructurado). Soportan formatos de calendario y cálculos con TimeDeltas.\n",
      "* **Cadenas:** `replace` (saneamiento y sustitución de texto) y `randomString` (generación determinista de claves/tokens aleatorios).\n",
      "\n",
      "### Especificación BNF (Sección VI)\n",
      "\n",
      "\n",
      "\n",
      "/* [CORRECCIÓN] Todas las subreglas de <util_command> están ahora completamente expandidas. */\n",
      "<util_command>    ::= <json_list_cmd> | <crypto_cmd> | <regex_cmd> | <datetime_cmd> | <stamp_cmd> | <string_cmd> | <replace_cmd>\n",
      "\n",
      "/* Manipulación de listas y JSON */\n",
      "<json_list_cmd>   ::= \"variableToList(\" <expression> \",\" <identifier> \")\"\n",
      "                    | \"itemFromList(\" <identifier> \",\" <expression> \",\" <identifier> \")\"\n",
      "                    | \"variableFromJSON(\" <identifier> \",\" <expression> \",\" <identifier> \")\"\n",
      "                    | \"AddVariableToJSON(\" <expression> \",\" <expression> \",\" <identifier> \")\"\n",
      "\n",
      "/* Criptografía */\n",
      "<crypto_cmd>      ::= \"encodeSHA256(\" <identifier_or_string> \",\" <identifier> \")\"\n",
      "                    | \"encodeMD5(\" <identifier_or_string> \",\" <identifier> \")\"\n",
      "\n",
      "/* Expresiones regulares */\n",
      "<regex_cmd>       ::= \"getRegex(\" <identifier> \",\" <stringliteral> \",\" <identifier> \")\"\n",
      "\n",
      "<datetime_cmd>    ::= \"getDateTime(\" <stringliteral> \",\" <expression> \",\" <stringliteral> \",\" <identifier> \")\"\n",
      "/*  Argumentos: formato_salida, epoch_origen, zona_horaria, destino */\n",
      "\n",
      "<stamp_cmd>       ::= \"stampToDatetime(\" <expression> \",\" <stringliteral> \",\" <expression> \",\" <identifier> \")\"\n",
      "/*  Argumentos: epoch_origen, formato, timedelta, destino */\n",
      "                    | \"getTimeStamp(\" <stringliteral> \",\" <stringliteral> \",\" <expression> \",\" <identifier> \")\"\n",
      "/*  Argumentos: fecha_string, formato_entrada, timedelta, destino */\n",
      "\n",
      "<string_cmd>      ::= \"randomString(\" <expression> \",\" <identifier> \")\"\n",
      "/*  Argumentos: longitud, destino */\n",
      "\n",
      "<replace_cmd>     ::= \"replace(\" <identifier_or_string> \",\" <stringliteral> \",\" <stringliteral> \",\" <identifier> \")\"\n",
      "/*  Argumentos: origen, patron_busqueda, reemplazo, destino */\n",
      "\n",
      "[3] id=chunk-3 source=Untitled\n",
      "\n",
      "\n",
      "---\n",
      "\n",
      "## SECCIÓN IX: Expresiones y Gramática Léxica Estricta\n",
      "\n",
      "Esta sección es el corazón matemático evaluador de AVAP. Define la jerarquía exacta (Precedencia) y provee soporte nativo para características avanzadas similares a Python.\n",
      "\n",
      "### 9.1 Cast de Tipos Explícito\n",
      "AVAP permite conversiones de tipos (Type Casting) en cualquier evaluación utilizando funciones constructoras estándar. Puedes transformar variables dinámicamente usando `int(var)`, `float(var)` o `str(var)`.\n",
      "\n",
      "### 9.2 Slicing y Comprensiones (Comprehensions)\n",
      "* **Slicing (Cortes):** Puedes extraer fragmentos de listas o strings utilizando la notación de dos puntos. Ejemplo: `mi_lista[1:4]` (extrae desde el índice 1 hasta el 3).\n",
      "* **Comprehensions:** AVAP soporta la construcción rápida de listas mediante iteradores en una sola línea, permitiendo filtrar y mapear colecciones enteras (ej. `[x * 2 for x in valores if x > 0]`).\n",
      "\n",
      "### 9.3 Análisis Léxico (Lexer) y Documentación\n",
      "AVAP cuenta con tres niveles de descarte de texto para anotaciones humanas:\n",
      "1. **Comentarios de Línea (`//`):** Ignora el texto hasta el salto de línea.\n",
      "2. **Comentarios de Bloque (`/* ... */`):** Para aislar bloques enteros multilínea.\n",
      "3. **Comentarios de Documentación (`///`):** Utilizados por analizadores de código o IDEs para generar documentación técnica automática (Docstrings) a partir del código fuente.\n",
      "\n",
      "### Especificación BNF (Sección IX)\n",
      "\n",
      "\n",
      "\n",
      "/* Jerarquía de Expresiones (Precedencia de menor a mayor) */\n",
      "<expression>       ::= <logical_or>\n",
      "<logical_or>       ::= <logical_and> ( \"or\" <logical_and> )*\n",
      "<logical_and>      ::= <logical_not> ( \"and\" <logical_not> )*\n",
      "<logical_not>      ::= \"not\" <logical_not> | <comparison>\n",
      "\n",
      "<comparison>       ::= <arithmetic> ( <comp_op> <arithmetic> )*\n",
      "<comp_op>          ::= \"==\" | \"!=\" | \"<\" | \">\" | \"<=\" | \">=\" | \"in\" | \"is\"\n",
      "\n",
      "<arithmetic>       ::= <term> ( ( \"+\" | \"-\" ) <term> )*\n",
      "<term>             ::= <factor> ( ( \"*\" | \"/\" | \"%\" ) <factor> )*\n",
      "<factor>           ::= ( \"+\" | \"-\" ) <factor> | <power>\n",
      "<power>            ::= <primary> [ \"**\" <factor> ]\n",
      "\n",
      "/* Primarios y Átomos (Accesos, Castings, Slicing, Métodos y Funciones)\n",
      "   La regla <primary> cubre también el acceso a métodos de objetos conector\n",
      "   (conector.metodo(...)) y el acceso por clave a sus resultados (resultado[\"key\"]) */\n",
      "<primary>          ::= <atom>\n",
      "                     | <primary> \".\" <identifier>\n",
      "                     | <primary> \"[\" <expression> \"]\"\n",
      "                     | <primary> \"[\" [<expression>] \":\" [<expression>] [\":\" [<expression>]] \"]\"\n",
      "                     | <primary> \"(\" [<argument_list>] \")\"\n",
      "\n",
      "<atom>             ::= <identifier>\n",
      "                     | \"$\" <identifier>\n",
      "                     | <literal>\n",
      "                     | \"(\" <expression> \")\"\n",
      "                     | <list_display>\n",
      "                     | <dict_display>\n",
      "\n",
      "/* Estructuras de Datos, Comprensiones y Argumentos */\n",
      "<list_display>     ::= \"[\" [<argument_list>] \"]\"\n",
      "                     | \"[\" <expression> \"for\" <identifier> \"in\" <expression> [<if_clause>] \"]\"\n",
      "<if_clause>        ::= \"if\" <expression>\n",
      "<dict_display>     ::= \"{\" [<key_datum_list>] \"}\"\n",
      "<key_datum_list>   ::= <key_datum> ( \",\" <key_datum> )*\n",
      "<key_datum>        ::= <expression> \":\" <expression>\n",
      "<argument_list>    ::= <expression> ( \",\" <expression> )*\n",
      "\n",
      "/* Tipo numérico unificado */\n",
      "<number>           ::= <floatnumber> | <integer>\n",
      "\n",
      "/* Literales (Tipos de Datos Primitivos Soportados) */\n",
      "<literal>          ::= <stringliteral> | <number> | <boolean> | \"None\"\n",
      "<boolean>          ::= \"True\" | \"False\"\n",
      "<integer>          ::= [0-9]+\n",
      "<floatnumber>      ::= [0-9]+ \".\" [0-9]* | \".\" [0-9]+\n",
      "\n",
      "/* Cadenas de Texto con soporte de secuencias de escape */\n",
      "<stringliteral>    ::= \"\\\"\" <text_double> \"\\\"\" | \"'\" <text_single> \"'\"\n",
      "<escape_sequence>  ::= \"\\\\\" ( \"\\\"\" | \"'\" | \"\\\\\" | \"n\" | \"t\" | \"r\" | \"0\" )\n",
      "<text_double>      ::= ( [^\"\\\\] | <escape_sequence> )*\n",
      "<text_single>      ::= ( [^'\\\\] | <escape_sequence> )*\n",
      "<identifier_or_string> ::= <identifier> | <stringliteral>\n",
      "\n",
      "/* Reglas de Comentarios para el Lexer\n",
      "   El lexer aplica longest-match: /// debe evaluarse ANTES que // */\n",
      "<doc_comment>      ::= \"///\" <any_text>\n",
      "<line_comment>     ::= \"//\" <any_text>\n",
      "<block_comment>    ::= \"/*\" <any_content> \"*/\"\n",
      "<any_text>         ::= [^\\r\\n]*\n",
      "<any_content>      ::= /* Cualquier secuencia de caracteres que no contenga la subcadena \"*/\" */\n",
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "What types of includes does AVAP have?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "AVAP has two main types of include:\n",
      "\n",
      "1. **<connector_instantiation>:** This is used to instantiate a connector, which could be for a database connection or a third-party API.\n",
      "2. **<orm_connector_init>:** This term seems to be related to initializing an Object-Relational Mapping (ORM) connector, indicating that the context suggests it's part of a specific ORM setup.\n",
      "\n",
      "Both types are treated similarly in terms of grammar but differ semantically by their purpose - one is for database connections or third-party APIs, while the other is specifically for connecting to ORMs. The engine selects the appropriate adapter based on runtime context.\n"
     ]
    }
   ],
   "source": [
    "a = stream_graph_updates(user_input, guided_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6b4da6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = agentic_graph.invoke({\"messages\": [{\"role\": \"user\", \"content\": user_input}]})\n",
    "print(\"Final result:\")\n",
    "result[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2342b1f1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "3707574b",
   "metadata": {},
   "source": [
    "### MTEB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d9657ec4",
   "metadata": {},
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "invalid syntax (2793878467.py, line 12)",
     "output_type": "error",
     "traceback": [
      "  \u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[18]\u001b[39m\u001b[32m, line 12\u001b[39m\n\u001b[31m    \u001b[39m\u001b[31mnorms = np.linalg.norm(x, axis=1, keepdims=True)agent_graph\u001b[39m\n                                                    ^\n\u001b[31mSyntaxError\u001b[39m\u001b[31m:\u001b[39m invalid syntax\n"
     ]
    }
   ],
   "source": [
    "from dataclasses import dataclass\n",
    "from typing import Any, Iterable\n",
    " \n",
    "import numpy as np\n",
    " \n",
    "import mteb\n",
    "from mteb.types import Array\n",
    "from mteb.models import SearchEncoderWrapper\n",
    " \n",
    " \n",
    "def _l2_normalize(x: np.ndarray, eps: float = 1e-12) -> np.ndarray:\n",
    "    norms = np.linalg.norm(x, axis=1, keepdims=True)agent_graph\n",
    "    return x / np.clip(norms, eps, None)\n",
    " \n",
    " \n",
    "def _to_text_list(batch: dict[str, Any]) -> list[str]:\n",
    "    \"\"\"\n",
    "    MTEB batched inputs can be:\n",
    "      - TextInput: {\"text\": [..]}\n",
    "      - CorpusInput: {\"title\": [..], \"body\": [..], \"text\": [..]}\n",
    "      - QueryInput: {\"query\": [..], \"instruction\": [..], \"text\": [..]}\n",
    "    We prefer \"text\" if present; otherwise compose from title/body or query/instruction.\n",
    "    \"\"\"\n",
    "    if \"text\" in batch and batch[\"text\"] is not None:\n",
    "        return list(batch[\"text\"])\n",
    " \n",
    "    if \"title\" in batch and \"body\" in batch:\n",
    "        titles = batch[\"title\"] or [\"\"] * len(batch[\"body\"])\n",
    "        bodies = batch[\"body\"] or [\"\"] * len(batch[\"title\"])\n",
    "        return [f\"{t} {b}\".strip() for t, b in zip(titles, bodies)]\n",
    " \n",
    "    if \"query\" in batch:\n",
    "        queries = list(batch[\"query\"])\n",
    "        instructions = batch.get(\"instruction\")\n",
    "        if instructions:\n",
    "            return [f\"{i} {q}\".strip() for q, i in zip(queries, instructions)]\n",
    "        return queries\n",
    " \n",
    "    raise ValueError(f\"Unsupported batch keys: {sorted(batch.keys())}\")\n",
    " \n",
    " \n",
    "@dataclass\n",
    "class OllamaLangChainEncoder:\n",
    "    lc_embeddings: Any  # OllamaEmbeddings implements embed_documents()\n",
    "    normalize: bool = True\n",
    " \n",
    "    # Optional metadata hook used by some wrappers; safe to keep as None for local runs\n",
    "    mteb_model_meta: Any = None\n",
    " \n",
    "    def encode(\n",
    "        self,\n",
    "        inputs: Iterable[dict[str, Any]],\n",
    "        *,\n",
    "        task_metadata: Any,\n",
    "        hf_split: str,\n",
    "        hf_subset: str,\n",
    "        prompt_type: Any = None,\n",
    "        **kwargs: Any,\n",
    "    ) -> Array:\n",
    "        all_vecs: list[np.ndarray] = []\n",
    " \n",
    "        for batch in inputs:\n",
    "            texts = _to_text_list(batch)\n",
    "            vecs = self.lc_embeddings.embed_documents(texts)\n",
    "            arr = np.asarray(vecs, dtype=np.float32)\n",
    "            if self.normalize:\n",
    "                arr = _l2_normalize(arr)\n",
    "            all_vecs.append(arr)\n",
    " \n",
    "        if not all_vecs:\n",
    "            return np.zeros((0, 0), dtype=np.float32)\n",
    " \n",
    "        return np.vstack(all_vecs)\n",
    " \n",
    "    def similarity(self, embeddings1: Array, embeddings2: Array) -> Array:\n",
    "        a = np.asarray(embeddings1, dtype=np.float32)\n",
    "        b = np.asarray(embeddings2, dtype=np.float32)\n",
    "        if self.normalize:\n",
    "            # dot == cosine if already normalized\n",
    "            return a @ b.T\n",
    "        a = _l2_normalize(a)\n",
    "        b = _l2_normalize(b)\n",
    "        return a @ b.T\n",
    " \n",
    "    def similarity_pairwise(self, embeddings1: Array, embeddings2: Array) -> Array:\n",
    "        a = np.asarray(embeddings1, dtype=np.float32)\n",
    "        b = np.asarray(embeddings2, dtype=np.float32)\n",
    "        if not self.normalize:\n",
    "            a = _l2_normalize(a)\n",
    "            b = _l2_normalize(b)\n",
    "        return np.sum(a * b, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85727a68",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder = OllamaLangChainEncoder(lc_embeddings=embeddings, normalize=True)\n",
    "search_model = SearchEncoderWrapper(encoder)\n",
    " \n",
    "tasks = mteb.get_tasks([\n",
    "    \"CodeSearchNetRetrieval\",\n",
    "    \"CodeSearchNetCCRetrieval\",\n",
    "    \"AppsRetrieval\",\n",
    "    \"StackOverflowDupQuestions\",\n",
    "])\n",
    "results = mteb.evaluate(\n",
    "    model=search_model,\n",
    "    tasks=tasks,\n",
    "    encode_kwargs={\"batch_size\": 32, \"show_progress_bar\": True}\n",
    ")\n",
    " \n",
    "print(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4052f229",
   "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.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
