LlamaIndex Integration

LlamaIndex defines BaseMemory and BaseRetriever contracts. CortexDB plugs into both via the v1 SDK.

Install

pip install cortexdbai llama-index llama-index-llms-openai

Memory

import os
from datetime import datetime, timezone
from uuid import uuid4
from typing import List
from llama_index.core.memory import BaseMemory
from llama_index.core.llms import ChatMessage, MessageRole
from cortexdb.v1 import V1Client

client = V1Client(
    api_url="https://api-v1.cortexdb.ai",
    actor="user:alice",
    bearer=os.environ["CORTEX_TOKEN"],
)
SCOPE = "org:acme/user:alice"


class CortexMemory(BaseMemory):
    def get(self, input: str | None = None) -> List[ChatMessage]:
        pack = client.recall(
            scope=SCOPE, view="holistic", query=input or "",
            include=["beliefs", "facts", "episodes"],
            budgets={"max_tokens": 3000},
        )
        ctx = pack["context_block"]
        if not ctx:
            return []
        return [ChatMessage(role=MessageRole.SYSTEM, content=f"<recall>{ctx}</recall>")]

    def get_all(self) -> List[ChatMessage]:
        return self.get(None)

    def put(self, message: ChatMessage) -> None:
        client.experience(
            scope=SCOPE, text=str(message.content),
            role=("assistant" if message.role == MessageRole.ASSISTANT else "user"),
            observed_at=datetime.now(timezone.utc).isoformat(),
            idempotency_key=f"li-{uuid4()}",
        )

    def set(self, messages: List[ChatMessage]) -> None:
        for m in messages:
            self.put(m)

    def reset(self) -> None:
        pass

from llama_index.core.agent import ReActAgent
from llama_index.llms.openai import OpenAI

agent = ReActAgent.from_tools(tools, llm=OpenAI(model="gpt-4o"), memory=CortexMemory())
agent.chat("Our API rate limit is 1000 req/s per tenant.")
agent.chat("What is our API rate limit?")   # recalls from CortexDB

Retriever

from llama_index.core.retrievers import BaseRetriever
from llama_index.core.schema import NodeWithScore, TextNode

class CortexRetriever(BaseRetriever):
    def __init__(self, scope: str, k: int = 8):
        super().__init__()
        self._scope = scope
        self._k = k

    def _retrieve(self, query) -> List[NodeWithScore]:
        pack = client.recall(scope=self._scope, query=str(query), include=["facts", "episodes"])
        nodes = []
        for fact in pack["layers"].get("facts", []):
            nodes.append(NodeWithScore(
                node=TextNode(
                    text=f"{fact['subject']['name']} {fact['predicate']} {fact['object']['value']}",
                    metadata={"layer": "fact", "id": fact["fact_id"]},
                ),
                score=fact["confidence"],
            ))
        return nodes[: self._k]

See also

• Python SDK
• Layers