Working refactor. Dependency clean-up pending.

private_gpt/server/chat/chat_router.py

@@ -1,5 +1,5 @@
 from fastapi import APIRouter, Depends, Request
-from llama_index.llms import ChatMessage, MessageRole
+from llama_index.core.llms import ChatMessage, MessageRole
 from pydantic import BaseModel
 from starlette.responses import StreamingResponse
 

private_gpt/server/chat/chat_service.py

@@ -1,14 +1,15 @@
 from dataclasses import dataclass
 
 from injector import inject, singleton
-from llama_index import ServiceContext, StorageContext, VectorStoreIndex
-from llama_index.chat_engine import ContextChatEngine, SimpleChatEngine
-from llama_index.chat_engine.types import (
+from llama_index.core.chat_engine import ContextChatEngine, SimpleChatEngine
+from llama_index.core.chat_engine.types import (
     BaseChatEngine,
 )
-from llama_index.indices.postprocessor import MetadataReplacementPostProcessor
-from llama_index.llms import ChatMessage, MessageRole
-from llama_index.types import TokenGen
+from llama_index.core.indices import VectorStoreIndex
+from llama_index.core.indices.postprocessor import MetadataReplacementPostProcessor
+from llama_index.core.llms import ChatMessage, MessageRole
+from llama_index.core.storage import StorageContext
+from llama_index.core.types import TokenGen
 from pydantic import BaseModel
 
 from private_gpt.components.embedding.embedding_component import EmbeddingComponent
@@ -75,20 +76,19 @@ class ChatService:
         embedding_component: EmbeddingComponent,
         node_store_component: NodeStoreComponent,
     ) -> None:
-        self.llm_service = llm_component
+        self.llm_component = llm_component
+        self.embedding_component = embedding_component
         self.vector_store_component = vector_store_component
         self.storage_context = StorageContext.from_defaults(
             vector_store=vector_store_component.vector_store,
             docstore=node_store_component.doc_store,
             index_store=node_store_component.index_store,
         )
-        self.service_context = ServiceContext.from_defaults(
-            llm=llm_component.llm, embed_model=embedding_component.embedding_model
-        )
         self.index = VectorStoreIndex.from_vector_store(
             vector_store_component.vector_store,
             storage_context=self.storage_context,
-            service_context=self.service_context,
+            llm=llm_component.llm,
+            embed_model=embedding_component.embedding_model,
             show_progress=True,
         )
 
@@ -102,10 +102,17 @@ class ChatService:
             vector_index_retriever = self.vector_store_component.get_retriever(
                 index=self.index, context_filter=context_filter
             )
+            # TODO ContextChatEngine is still not migrated by LlamaIndex to accept
+            # llm directly, so we are passing legacy ServiceContext until it is fixed.
+            from llama_index.core import ServiceContext
             return ContextChatEngine.from_defaults(
                 system_prompt=system_prompt,
                 retriever=vector_index_retriever,
-                service_context=self.service_context,
+                llm=self.llm_component.llm,  # Takes no effect at the moment
+                service_context=ServiceContext.from_defaults(
+                    llm=self.llm_component.llm,
+                    embed_model=self.embedding_component.embedding_model,
+                ),
                 node_postprocessors=[
                     MetadataReplacementPostProcessor(target_metadata_key="window"),
                 ],
@@ -113,7 +120,7 @@ class ChatService:
         else:
             return SimpleChatEngine.from_defaults(
                 system_prompt=system_prompt,
-                service_context=self.service_context,
+                llm=self.llm_component.llm,
             )
 
     def stream_chat(

private_gpt/server/chunks/chunks_service.py

@@ -1,8 +1,9 @@
 from typing import TYPE_CHECKING, Literal
 
 from injector import inject, singleton
-from llama_index import ServiceContext, StorageContext, VectorStoreIndex
-from llama_index.schema import NodeWithScore
+from llama_index.core.indices import VectorStoreIndex
+from llama_index.core.schema import NodeWithScore
+from llama_index.core.storage import StorageContext
 from pydantic import BaseModel, Field
 
 from private_gpt.components.embedding.embedding_component import EmbeddingComponent
@@ -15,7 +16,7 @@ from private_gpt.open_ai.extensions.context_filter import ContextFilter
 from private_gpt.server.ingest.model import IngestedDoc
 
 if TYPE_CHECKING:
-    from llama_index.schema import RelatedNodeInfo
+    from llama_index.core.schema import RelatedNodeInfo
 
 
 class Chunk(BaseModel):
@@ -63,14 +64,13 @@ class ChunksService:
         node_store_component: NodeStoreComponent,
     ) -> None:
         self.vector_store_component = vector_store_component
+        self.llm_component = llm_component
+        self.embedding_component = embedding_component
         self.storage_context = StorageContext.from_defaults(
             vector_store=vector_store_component.vector_store,
             docstore=node_store_component.doc_store,
             index_store=node_store_component.index_store,
         )
-        self.query_service_context = ServiceContext.from_defaults(
-            llm=llm_component.llm, embed_model=embedding_component.embedding_model
-        )
 
     def _get_sibling_nodes_text(
         self, node_with_score: NodeWithScore, related_number: int, forward: bool = True
@@ -103,7 +103,8 @@ class ChunksService:
         index = VectorStoreIndex.from_vector_store(
             self.vector_store_component.vector_store,
             storage_context=self.storage_context,
-            service_context=self.query_service_context,
+            llm=self.llm_component.llm,
+            embed_model=self.embedding_component.embedding_model,
             show_progress=True,
         )
         vector_index_retriever = self.vector_store_component.get_retriever(

private_gpt/server/ingest/ingest_service.py

@@ -4,11 +4,8 @@ from pathlib import Path
 from typing import AnyStr, BinaryIO
 
 from injector import inject, singleton
-from llama_index import (
-    ServiceContext,
-    StorageContext,
-)
-from llama_index.node_parser import SentenceWindowNodeParser
+from llama_index.core.node_parser import SentenceWindowNodeParser
+from llama_index.core.storage import StorageContext
 
 from private_gpt.components.embedding.embedding_component import EmbeddingComponent
 from private_gpt.components.ingest.ingest_component import get_ingestion_component
@@ -40,17 +37,12 @@ class IngestService:
             index_store=node_store_component.index_store,
         )
         node_parser = SentenceWindowNodeParser.from_defaults()
-        self.ingest_service_context = ServiceContext.from_defaults(
-            llm=self.llm_service.llm,
-            embed_model=embedding_component.embedding_model,
-            node_parser=node_parser,
-            # Embeddings done early in the pipeline of node transformations, right
-            # after the node parsing
-            transformations=[node_parser, embedding_component.embedding_model],
-        )
 
         self.ingest_component = get_ingestion_component(
-            self.storage_context, self.ingest_service_context, settings=settings()
+            self.storage_context,
+            embed_model=embedding_component.embedding_model,
+            transformations=[node_parser, embedding_component.embedding_model],
+            settings=settings(),
         )
 
     def _ingest_data(self, file_name: str, file_data: AnyStr) -> list[IngestedDoc]:

private_gpt/server/ingest/model.py

@@ -1,6 +1,6 @@
 from typing import Any, Literal
 
-from llama_index import Document
+from llama_index.core.schema import Document
 from pydantic import BaseModel, Field