Update Llama Stack to 0.2.22

Author: syedriko

pyproject.toml

@@ -20,6 +20,9 @@ line-length = 100
 [tool.mypy]
 disable_error_code = ["union-attr", "return-value", "arg-type", "import-untyped"]
 ignore_missing_imports = true
+mypy_path = ["src"]
+namespace_packages = true
+explicit_package_bases = true
 
 [tool.pdm]
 distribution = true
@@ -34,19 +37,19 @@ version = "0.1.0"
 description = "RAG content for OpenShift LightSpeed."
 authors = []
 dependencies = [
-    "PyYAML==6.0.2",
+    "PyYAML>=6.0.2",
     "huggingface_hub>=0.33.4",
-    "llama-index==0.12.51",
-    "llama-index-vector-stores-faiss==0.3.0",
-    "llama-index-embeddings-huggingface==0.4.0",
-    "llama-index-readers-file==0.4.11",
-    "faiss-cpu==1.11.0.post1",
+    "llama-index>=0.12.51",
+    "llama-index-vector-stores-faiss>=0.3.0",
+    "llama-index-embeddings-huggingface>=0.4.0",
+    "llama-index-readers-file>=0.4.11",
+    "faiss-cpu>=1.11.0.post1",
     "llama-index-vector-stores-postgres>=0.5.4",
-    "torch==2.7.1",
-    "llama-stack==0.2.16",
-    "llama-stack-client==0.2.16",
-    "aiosqlite==0.21.0",
-    "sqlite-vec==0.1.6",
+    "torch>=2.7.1",
+    "llama-stack==0.2.22",
+    "llama-stack-client==0.2.22",
+    "aiosqlite>=0.21.0",
+    "sqlite-vec>=0.1.6",
     "tomlkit",
 ]
 requires-python = "==3.12.*"

scripts/query_rag.py

@@ -12,6 +12,7 @@
 import yaml
 from llama_index.core import Settings, load_index_from_storage
 from llama_index.core.llms.utils import resolve_llm
+from llama_index.core.schema import NodeWithScore, TextNode
 from llama_index.core.storage.storage_context import StorageContext
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from llama_index.vector_stores.faiss import FaissVectorStore
@@ -35,20 +36,34 @@ def _llama_index_query(args: argparse.Namespace) -> None:
 
     if args.node is not None:
         node = storage_context.docstore.get_node(args.node)
-        result = {
-            "query": args.query,
-            "type": "single_node",
-            "node_id": args.node,
-            "node": {
-                "id": node.node_id,
-                "text": node.text,
-                "metadata": node.metadata if hasattr(node, "metadata") else {},
-            },
-        }
-        if args.json:
-            print(json.dumps(result, indent=2))
+        if isinstance(node, TextNode):
+            result = {
+                "query": args.query,
+                "type": "single_node",
+                "node_id": args.node,
+                "node": {
+                    "id": node.node_id,
+                    "text": node.text,
+                    "metadata": node.metadata if hasattr(node, "metadata") else {},
+                },
+            }
+            if args.json:
+                print(json.dumps(result, indent=2))
+            else:
+                print(node)
         else:
-            print(node)
+            logging.warning(
+                f"Node {args.node} is not a TextNode, type: {type(node).__name__}"
+            )
+            if args.json:
+                result = {
+                    "query": args.query,
+                    "type": "single_node",
+                    "node_id": args.node,
+                    "error": f"Node is not a TextNode (type: {type(node).__name__})",
+                }
+                print(json.dumps(result, indent=2))
+            exit(1)
     else:
         retriever = vector_index.as_retriever(similarity_top_k=args.top_k)
         nodes = retriever.retrieve(args.query)
@@ -88,13 +103,18 @@ def _llama_index_query(args: argparse.Namespace) -> None:
             "nodes": [],
         }
         for node in nodes:  # type: ignore
-            node_data = {
-                "id": node.node_id,
-                "score": node.score,
-                "text": node.text,
-                "metadata": node.metadata if hasattr(node, "metadata") else {},
-            }
-            result["nodes"].append(node_data)
+            if isinstance(node, NodeWithScore):
+                node_data = {
+                    "id": node.node_id,
+                    "score": node.score,
+                    "text": node.text,
+                    "metadata": node.metadata if hasattr(node, "metadata") else {},
+                }
+                result["nodes"].append(node_data)
+            else:
+                logging.debug(
+                    f"Skipping node of type {type(node).__name__}, expected NodeWithScore"
+                )
 
         if args.json:
             print(json.dumps(result, indent=2))
@@ -134,7 +154,7 @@ def _llama_stack_query(args: argparse.Namespace) -> None:
     yaml.safe_dump(cfg, open(cfg_file, "w", encoding="utf-8"))
 
     stack_lib = importlib.import_module("llama_stack")
-    client = stack_lib.distribution.library_client.LlamaStackAsLibraryClient(cfg_file)
+    client = stack_lib.core.library_client.LlamaStackAsLibraryClient(cfg_file)
     client.initialize()
 
     # No need to register the DB as it's defined in llama-stack.yaml

src/lightspeed_rag_content/document_processor.py

@@ -280,12 +280,10 @@ def __init__(self, config: _Config):
         # Not using importlib to help with typechecking
         import llama_stack  # pylint: disable=C0415
 
-        self.document_class = llama_stack.apis.tools.rag_tool.RAGDocument  # type: ignore
-        self.client_class = (
-            llama_stack.distribution.library_client.LlamaStackAsLibraryClient  # type: ignore
-        )
+        self.document_class = llama_stack.apis.tools.rag_tool.RAGDocument
+        self.client_class = llama_stack.core.library_client.LlamaStackAsLibraryClient
         self.documents: list[
-            dict[str, Any] | llama_stack.apis.tools.rag_tool.RAGDocument  # type: ignore
+            dict[str, Any] | llama_stack.apis.tools.rag_tool.RAGDocument
         ] = []
 
     def write_yaml_config(self, index_id: str, filename: str, db_file: str) -> None:
@@ -313,7 +311,7 @@ def _start_llama_stack(self, cfg_file: str) -> Any:
         """Start llama-stack as a library and return the client.
 
         Return type is really
-          llama_stack.distribution.library_client.LlamaStackAsLibraryClient
+          llama_stack.core.library_client.LlamaStackAsLibraryClient
 
         But we do dynamic import, so we don't have it for static typechecking
         """