LlamaIndex: Interactive Notes

"""Key-free stand-ins for LlamaIndex. Real code: Settings.llm = OpenAI(model="gpt-4o-mini"); Settings.embed_model = OpenAIEmbedding()."""
import re, math
from llama_index.core.embeddings import BaseEmbedding

class KeywordEmbedding(BaseEmbedding):
    def _vec(self, text, dim=256):
        v = [0.0] * dim
        for w in re.findall(r"[a-z]+", text.lower()): v[sum(map(ord, w[:5])) % dim] += 1
        n = math.sqrt(sum(x * x for x in v)) or 1
        return [x / n for x in v]
    def _get_text_embedding(self, text): return self._vec(text)
    def _get_query_embedding(self, query): return self._vec(query)
    async def _aget_query_embedding(self, query): return self._vec(query)

# pip install llama-index-core   (plus llama-index-llms-openai etc. for real models)
from llama_index.core import Document, VectorStoreIndex, Settings
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.llms import MockLLM
from li_fake import KeywordEmbedding

Settings.embed_model = KeywordEmbedding()        # real: OpenAIEmbedding()
Settings.llm = MockLLM(max_tokens=12)            # real: OpenAI(model="gpt-4o-mini")

docs = [Document(text="Premature withdrawal of an FD reduces the interest rate by 1%. "
                      "Senior citizens get an extra 0.50% on all FD tenures.", metadata={"file": "fd_policy.pdf"}),
        Document(text="A home loan EMI is due on the 5th of every month. Late payment costs ₹500.",
                 metadata={"file": "loan_policy.pdf"})]

nodes = SentenceSplitter(chunk_size=25, chunk_overlap=0).get_nodes_from_documents(docs)   # Documents -> Nodes
print(len(nodes), "nodes; first:", nodes[0].text, "|", nodes[0].metadata)

index = VectorStoreIndex(nodes)                                     # embed + store
retriever = index.as_retriever(similarity_top_k=1)
for hit in retriever.retrieve("What is the penalty for withdrawing an FD early?"):
    print(f"retrieved ({hit.score:.2f}):", hit.node.text, "<-", hit.node.metadata["file"])

engine = index.as_query_engine(similarity_top_k=1)                  # retrieve + prompt + LLM in one call
resp = engine.query("When is the EMI due?")
print("sources used:", [n.node.metadata["file"] for n in resp.source_nodes])

from llama_index.core import Settings, StorageContext, load_index_from_storage
from llama_index.llms.openai import OpenAI                 # pip install llama-index-llms-openai
from llama_index.embeddings.openai import OpenAIEmbedding  # pip install llama-index-embeddings-openai

Settings.llm = OpenAI(model="gpt-4o-mini")
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")

index.storage_context.persist("./storage")                 # save once ...
index = load_index_from_storage(StorageContext.from_defaults(persist_dir="./storage"))   # ... reload later

import asyncio
from llama_index.core.workflow import Workflow, step, Event, StartEvent, StopEvent

class Classified(Event):
    category: str
    text: str

class TicketFlow(Workflow):
    @step
    async def classify(self, ev: StartEvent) -> Classified:          # each step listens for an event type
        cat = "card" if "card" in ev.text.lower() else "general"
        return Classified(category=cat, text=ev.text)

    @step
    async def respond(self, ev: Classified) -> StopEvent:             # ...and emits the next event
        reply = {"card": "Your card is blocked; a new one is on its way.",
                 "general": "Thanks, an agent will reply within 24 hours."}[ev.category]
        return StopEvent(result=f"[{ev.category}] {reply}")

async def main():
    flow = TicketFlow(timeout=10)
    for text in ["My card was stolen", "I want to update my address"]:
        print(await flow.run(text=text))

asyncio.run(main())

from llama_index.core.agent.workflow import FunctionAgent

def fd_rate(years: int) -> str:
    """Return the FD interest rate for a tenure in years."""
    return {1: "6.8%", 3: "7.1%"}.get(years, "not offered")

agent = FunctionAgent(tools=[fd_rate], llm=Settings.llm,         # needs a real model that supports tool calling
                      system_prompt="You are NovaBank's assistant.")
# inside an async function:  response = await agent.run("What is the 3-year FD rate?")