feat: add retrieval functionality and update execution counts in notebooks
This commit is contained in:
pseco
parent
51488b3ee6
commit
4b0be0b80b
|
|
@ -0,0 +1,318 @@
|
|||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "9f97dd1e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Libraries"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "9e974df6",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from __future__ import annotations\n",
|
||||
"\n",
|
||||
"from typing import List, Dict, Any, Optional, TypedDict\n",
|
||||
"from pydantic import BaseModel, Field\n",
|
||||
"import os\n",
|
||||
"from elasticsearch import Elasticsearch\n",
|
||||
"from langchain_community.embeddings import OllamaEmbeddings\n",
|
||||
"from langchain_community.llms import Ollama\n",
|
||||
"from langchain_core.messages import HumanMessage, SystemMessage\n",
|
||||
"from langchain_elasticsearch import ElasticsearchStore\n",
|
||||
"from langgraph.graph import StateGraph, END"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "30edcecc",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"/tmp/ipykernel_837760/4144511709.py:6: LangChainDeprecationWarning: The class `OllamaEmbeddings` was deprecated in LangChain 0.3.1 and will be removed in 1.0.0. An updated version of the class exists in the `langchain-ollama package and should be used instead. To use it run `pip install -U `langchain-ollama` and import as `from `langchain_ollama import OllamaEmbeddings``.\n",
|
||||
" embeddings = OllamaEmbeddings(\n",
|
||||
"/tmp/ipykernel_837760/4144511709.py:8: LangChainDeprecationWarning: The class `Ollama` was deprecated in LangChain 0.3.1 and will be removed in 1.0.0. An updated version of the class exists in the `langchain-ollama package and should be used instead. To use it run `pip install -U `langchain-ollama` and import as `from `langchain_ollama import OllamaLLM``.\n",
|
||||
" llm = Ollama(base_url=base_url, model=model_name)\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"es = Elasticsearch(os.getenv(\"ELASTICSEARCH_LOCAL_URL\"))\n",
|
||||
"index_name = os.getenv(\"ELASTICSEARCH_INDEX\")\n",
|
||||
"base_url = os.getenv(\"LLM_BASE_LOCAL_URL\")\n",
|
||||
"model_name = os.getenv(\"OLLAMA_MODEL_NAME\")\n",
|
||||
"\n",
|
||||
"embeddings = OllamaEmbeddings(base_url=base_url, model=model_name)\n",
|
||||
"llm = Ollama(base_url=base_url, model=model_name)\n",
|
||||
"\n",
|
||||
"vector_store = ElasticsearchStore(\n",
|
||||
" es_url=es,\n",
|
||||
" index_name=index_name,\n",
|
||||
" embedding=embeddings,\n",
|
||||
" query_field=\"text\",\n",
|
||||
" vector_query_field=\"embedding\",\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "1e1dd107",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# State"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"id": "19e723e2",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"class RAGState(TypedDict, total=False):\n",
|
||||
" question: str\n",
|
||||
" query_embedding: List[float]\n",
|
||||
" docs: List[Dict[str, Any]]\n",
|
||||
" answer: str\n",
|
||||
" parsed: Dict[str, Any]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "90162558",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Schema"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"id": "acfe33d8",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"class AnswerSchema(BaseModel):\n",
|
||||
" answer: str = Field(..., description=\"Final answer to the user\")\n",
|
||||
" citations: List[str] = Field(default_factory=list, description=\"List of sources/ids used\")\n",
|
||||
" confidence: str = Field(..., description=\"low|medium|high\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "88a4aba1",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# ES Retrieval"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"id": "6c9e9b89",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def es_vector_search(\n",
|
||||
" es_client: Elasticsearch,\n",
|
||||
" index: str,\n",
|
||||
" query_vector: List[float],\n",
|
||||
" k: int = 5,\n",
|
||||
" num_candidates: int = 50,\n",
|
||||
" title_filter: Optional[str] = None,\n",
|
||||
"):\n",
|
||||
" \"\"\"\n",
|
||||
" Uses Elasticsearch kNN search against your field: 'vector'\n",
|
||||
" and returns _source: text + metadata.title (matches your mapping).\n",
|
||||
" \"\"\"\n",
|
||||
" body: Dict[str, Any] = {\n",
|
||||
" \"knn\": {\n",
|
||||
" \"field\": \"vector\", # <-- your mapping\n",
|
||||
" \"query_vector\": query_vector,\n",
|
||||
" \"k\": k,\n",
|
||||
" \"num_candidates\": num_candidates,\n",
|
||||
" },\n",
|
||||
" \"_source\": [\"text\", \"metadata.title\"], # <-- your mapping\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" # Optional: filter by title (exact via keyword subfield, or fuzzy via text)\n",
|
||||
" if title_filter:\n",
|
||||
" body[\"query\"] = {\n",
|
||||
" \"bool\": {\n",
|
||||
" \"filter\": [\n",
|
||||
" {\"term\": {\"metadata.title.keyword\": title_filter}}\n",
|
||||
" ]\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" res = es_client.search(index=index, body=body)\n",
|
||||
" return res[\"hits\"][\"hits\"]\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def normalize_hits(hits: List[Dict[str, Any]]) -> List[Dict[str, Any]]:\n",
|
||||
" docs: List[Dict[str, Any]] = []\n",
|
||||
" for h in hits:\n",
|
||||
" src = h.get(\"_source\", {}) or {}\n",
|
||||
" meta = src.get(\"metadata\", {}) or {}\n",
|
||||
" docs.append(\n",
|
||||
" {\n",
|
||||
" \"id\": h.get(\"_id\"),\n",
|
||||
" \"score\": h.get(\"_score\"),\n",
|
||||
" \"text\": src.get(\"text\", \"\"),\n",
|
||||
" \"title\": meta.get(\"title\", None),\n",
|
||||
" }\n",
|
||||
" )\n",
|
||||
" return docs\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def format_context(docs: List[Dict[str, Any]]) -> str:\n",
|
||||
" chunks = []\n",
|
||||
" for i, d in enumerate(docs, 1):\n",
|
||||
" title = d.get(\"title\") or \"Untitled\"\n",
|
||||
" doc_id = d.get(\"id\") or f\"chunk-{i}\"\n",
|
||||
" chunks.append(f\"[{i}] id={doc_id} title={title}\\n{d.get('text','')}\")\n",
|
||||
" return \"\\n\\n\".join(chunks)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "2591e778",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Nodes"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "69c41a89",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def embed_query(state: RAGState) -> RAGState:\n",
|
||||
" q = state[\"question\"]\n",
|
||||
" vec = embeddings.embed_query(q)\n",
|
||||
" return {\"query_embedding\": vec}\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def retrieve(state: RAGState) -> RAGState:\n",
|
||||
" base_retriever = vector_store.as_retriever(\n",
|
||||
" search_type=\"similarity\",\n",
|
||||
" search_kwargs={\"k\": 4}\n",
|
||||
" ) \n",
|
||||
" docs = base_retriever.invoke(\"What does the text say about agricultural techniques?\")\n",
|
||||
" docs = normalize_hits(hits)\n",
|
||||
" return {\"docs\": docs}\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def generate_answer(state: RAGState) -> RAGState:\n",
|
||||
" question = state[\"question\"]\n",
|
||||
" docs = state.get(\"docs\", [])\n",
|
||||
" context = format_context(docs)\n",
|
||||
"\n",
|
||||
" system = SystemMessage(\n",
|
||||
" content=(\n",
|
||||
" \"You are a helpful RAG assistant. Use ONLY the provided context. \"\n",
|
||||
" \"If the context is insufficient, say what is missing and ask a precise follow-up question. \"\n",
|
||||
" \"Cite sources like [1], [2] based on the context chunks.\"\n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
" user = HumanMessage(\n",
|
||||
" content=f\"Question:\\n{question}\\n\\nContext:\\n{context}\\n\\nWrite the best possible answer.\"\n",
|
||||
" )\n",
|
||||
" resp = llm.invoke([system, user])\n",
|
||||
" # Handle both string and message object responses\n",
|
||||
" answer = resp.content if hasattr(resp, 'content') else resp\n",
|
||||
" return {\"answer\": answer}"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f3933c13",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Graph"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "d6810eb8",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"graph = StateGraph(RAGState)\n",
|
||||
"\n",
|
||||
"graph.add_node(\"embed_query\", embed_query)\n",
|
||||
"graph.add_node(\"retrieve\", retrieve)\n",
|
||||
"graph.add_node(\"generate_answer\", generate_answer)\n",
|
||||
"\n",
|
||||
"graph.set_entry_point(\"embed_query\")\n",
|
||||
"graph.add_edge(\"embed_query\", \"retrieve\")\n",
|
||||
"graph.add_edge(\"retrieve\", \"generate_answer\")\n",
|
||||
"graph.add_edge(\"generate_answer\", END)\n",
|
||||
"\n",
|
||||
"dag_app = graph.compile()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "6528efac",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Retrieve"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"id": "568708ba",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"RAW ANSWER:\n",
|
||||
" New agricultural techniques include vertical farming, which integrates plants into multi-level structures in densely populated cities to optimize space and reduce transportation dependency. This approach uses hydroponic systems and automated nutrient control to grow food without traditional soil. Additionally, agriculture buildings are designed to seamlessly integrate with urban environments, addressing economic challenges while enhancing sustainability through the use of renewable energy sources.\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"out = dag_app.invoke({\"question\": \"What are new agricultural techniques about?\"})\n",
|
||||
"\n",
|
||||
"print(\"RAW ANSWER:\\n\", out[\"answer\"])"
|
||||
]
|
||||
}
|
||||
],
|
||||
"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
|
||||
}
|
||||
|
|
@ -380,7 +380,7 @@
|
|||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"execution_count": 18,
|
||||
"id": "1ed4c817",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
|
|
@ -459,6 +459,50 @@
|
|||
" print(\"Source:\", hit[\"_source\"])\n",
|
||||
" print(\"-\" * 40)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "d823650e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Retrive"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 27,
|
||||
"id": "5732a27d",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"base_retriever = db.as_retriever(\n",
|
||||
" search_type=\"similarity\",\n",
|
||||
" search_kwargs={\"k\": 1}\n",
|
||||
" ) \n",
|
||||
"\n",
|
||||
"docs = base_retriever.invoke(\"What does the text say about agricultural techniques?\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 28,
|
||||
"id": "28771235",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[Document(metadata={'title': 'Demo'}, page_content=' some cases, agricultural buildings are architecturally designed to integrate seamlessly into the urban environment. Although it still faces economic challenges, vertical farming represents a potential solution for food security in megacities.')]"
|
||||
]
|
||||
},
|
||||
"execution_count": 28,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"docs"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
|
|
|
|||
Loading…
Reference in New Issue