Amend the code for llama-index v0.10.x.

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
wxywb · Feb 26, 2024 · ceaefab · ceaefab
1 parent 506fd22
commit ceaefab
Show file tree

Hide file tree

Showing 7 changed files with 55 additions and 89 deletions.
diff --git a/custom/history_sentence_window.py b/custom/history_sentence_window.py
@@ -2,13 +2,14 @@
 from typing import Any, Callable, List, Optional, Sequence
 from bisect import bisect_right
 
-from llama_index.bridge.pydantic import Field
-from llama_index.callbacks.base import CallbackManager
-from llama_index.node_parser.interface import NodeParser
-from llama_index.node_parser.node_utils import build_nodes_from_splits
-from llama_index.node_parser.text.utils import split_by_sentence_tokenizer
-from llama_index.schema import BaseNode, Document, MetadataMode
-from llama_index.utils import get_tqdm_iterable
+from llama_index.core.bridge.pydantic import Field
+from llama_index.core.callbacks import CallbackManager
+from llama_index.core.node_parser import NodeParser
+from llama_index.core.node_parser.node_utils import build_nodes_from_splits
+from llama_index.core.node_parser.text.utils import split_by_sentence_tokenizer
+from llama_index.core.schema import BaseNode, MetadataMode
+from llama_index.core import Document
+from llama_index.core.utils import get_tqdm_iterable
 
 DEFAULT_WINDOW_SIZE = 3
 DEFAULT_WINDOW_METADATA_KEY = "window"

diff --git a/custom/llms/GeminiLLM.py b/custom/llms/GeminiLLM.py
@@ -4,22 +4,9 @@
     Sequence,
 )
 
-from llama_index.bridge.pydantic import Field
-from llama_index.llms.base import (
-    llm_chat_callback,
-    llm_completion_callback,
-)
-from llama_index.llms.llm import LLM
-from llama_index.llms.types import (
-    ChatMessage,
-    ChatResponse,
-    ChatResponseAsyncGen,
-    ChatResponseGen,
-    CompletionResponse,
-    CompletionResponseAsyncGen,
-    CompletionResponseGen,
-    LLMMetadata,
-)
+from llama_index.core.bridge.pydantic import Field
+from llama_index.core.llms import LLM
+from llama_index.core.llms import ChatMessage, ChatResponse, ChatResponseAsyncGen, ChatResponseGen, CompletionResponse, CompletionResponseAsyncGen, CompletionResponseGen, LLMMetadata
 
 DEFAULT_MODEL = "gemini-pro"
 
@@ -95,7 +82,6 @@ def call_with_prompt(self, prompt):
         if response is not None:
             return response.text
 
-    @llm_completion_callback()
     def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
         answer = self.call_with_prompt(prompt)
 
@@ -119,33 +105,26 @@ def metadata(self) -> LLMMetadata:
 
     # 下面是实现Interface必要的方法
     # 但这里用不到，所以都是pass
-    @llm_completion_callback()
     async def astream_complete(self) -> CompletionResponseAsyncGen:
         pass
 
     async def _astream_chat(self) -> ChatResponseAsyncGen:
         pass
 
-    @llm_chat_callback()
     async def astream_chat(self) -> ChatResponseAsyncGen:
         pass
 
-    @llm_chat_callback()
     def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
         pass
 
-    @llm_chat_callback()
     def stream_chat(self) -> ChatResponseGen:
         pass
 
-    @llm_completion_callback()
     def stream_complete(self) -> CompletionResponseGen:
         pass
 
-    @llm_chat_callback()
     async def achat(self) -> ChatResponse:
         pass
 
-    @llm_completion_callback()
     async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
-        pass
+        pass
diff --git a/custom/llms/QwenLLM.py b/custom/llms/QwenLLM.py
@@ -4,22 +4,9 @@
     Sequence,
 )
 
-from llama_index.bridge.pydantic import Field
-from llama_index.llms.base import (
-    llm_chat_callback,
-    llm_completion_callback,
-)
-from llama_index.llms.llm import LLM
-from llama_index.llms.types import (
-    ChatMessage,
-    ChatResponse,
-    ChatResponseAsyncGen,
-    ChatResponseGen,
-    CompletionResponse,
-    CompletionResponseAsyncGen,
-    CompletionResponseGen,
-    LLMMetadata,
-)
+from llama_index.core.bridge.pydantic import Field
+from llama_index.core.llms import LLM
+from llama_index.core.llms import ChatMessage, ChatResponse, ChatResponseAsyncGen, ChatResponseGen, CompletionResponse, CompletionResponseAsyncGen, CompletionResponseGen, LLMMetadata
 
 DEFAULT_MODEL = "qwen-max"
 #export DASHSCOPE_API_KEY="YOUR_KEY"
@@ -79,7 +66,6 @@ def call_with_prompt( self, prompt ):
 
             raise Exception( errMessage ) 
 
-    @llm_completion_callback()
     def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
         answer = self.call_with_prompt( prompt )
 
@@ -103,33 +89,26 @@ def metadata(self) -> LLMMetadata:
 
     # 下面是实现Interface必要的方法
     # 但这里用不到，所以都是pass
-    @llm_completion_callback()
     async def astream_complete() -> CompletionResponseAsyncGen:
         pass
 
     async def _astream_chat() -> ChatResponseAsyncGen:
         pass
 
-    @llm_chat_callback()
     async def astream_chat() -> ChatResponseAsyncGen:
         pass
 
-    @llm_chat_callback()
     def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
         pass
 
-    @llm_chat_callback()
     def stream_chat() -> ChatResponseGen:
         pass
 
-    @llm_completion_callback()
     def stream_complete() -> CompletionResponseGen:
         pass
 
-    @llm_chat_callback()
     async def achat() -> ChatResponse:
         pass
 
-    @llm_completion_callback()
     async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
-        pass
+        pass
diff --git a/custom/zilliz/base.py b/custom/zilliz/base.py
@@ -10,12 +10,14 @@
 
 import requests
 
-from llama_index.core.base_retriever import BaseRetriever
-from llama_index.data_structs.data_structs import IndexDict, IndexStructType
-from llama_index.indices.managed.base import BaseManagedIndex, IndexType
-from llama_index.schema import BaseNode, Document
-from llama_index.service_context import ServiceContext
-from llama_index.storage.storage_context import StorageContext
+from llama_index.core.retrievers import BaseRetriever
+from llama_index.core.data_structs import IndexDict
+from llama_index.core import IndexStructType
+from llama_index.core.indices.managed.base import BaseManagedIndex, IndexType
+from llama_index.core.schema import BaseNode
+from llama_index.core import Document
+from llama_index.core import ServiceContext
+from llama_index.core import StorageContext
 
 logger = logging.getLogger(__name__)
 
@@ -108,7 +110,7 @@ def insert_doc_url(self, url: str, metadata: Optional[Dict] = None) -> None:
 
 
         Example:
-        >>> from llama_index.indices import ZillizCloudPipelineIndex
+from llama_index.indices.managed.zilliz import ZillizCloudPipelineIndex
         >>> index = ZillizCloudPipelineIndex(
         >>>     project_id='YOUR_ZILLIZ_CLOUD_PROJECT_ID',
         >>>     cluster_id='YOUR_ZILLIZ_CLOUD_CLUSTER_ID',
@@ -212,7 +214,7 @@ def create_pipelines(
             A dictionary of pipeline ids for INGESTION, SEARCH, and DELETION pipelines.
 
         Example:
-            >>> from llama_index.indices import ZillizCloudPipelineIndex
+from llama_index.indices.managed.zilliz import ZillizCloudPipelineIndex
             >>> index = ZillizCloudPipelineIndex(
             >>>     project_id='YOUR_ZILLIZ_CLOUD_PROJECT_ID',
             >>>     cluster_id='YOUR_ZILLIZ_CLOUD_CLUSTER_ID',
@@ -336,7 +338,7 @@ def from_document_url(
             An initialized ZillizCloudPipelineIndex
 
         Example:
-            >>> from llama_index.indices import ZillizCloudPipelineIndex
+from llama_index.indices.managed.zilliz import ZillizCloudPipelineIndex
             >>> index = ZillizCloudPipelineIndex.from_document_url(
             >>>     url='https://oss_bucket.test_doc.ext',
             >>>     project_id='YOUR_ZILLIZ_CLOUD_PROJECT_ID',

diff --git a/custom/zilliz/retriever.py b/custom/zilliz/retriever.py
@@ -3,13 +3,14 @@
 
 import requests
 
-from llama_index.callbacks.base import CallbackManager
-from llama_index.constants import DEFAULT_SIMILARITY_TOP_K
-from llama_index.core.base_retriever import BaseRetriever
-from llama_index.indices.managed.zilliz.base import ZillizCloudPipelineIndex
-from llama_index.indices.query.schema import QueryBundle
-from llama_index.schema import NodeWithScore, QueryBundle, TextNode
-from llama_index.vector_stores.types import MetadataFilters
+from llama_index.core.callbacks import CallbackManager
+from llama_index.core.constants import DEFAULT_SIMILARITY_TOP_K
+from llama_index.core.retrievers import BaseRetriever
+from llama_index.indices.managed.zilliz import ZillizCloudPipelineIndex
+from llama_index.core import QueryBundle
+from llama_index.core.schema import NodeWithScore, TextNode
+from llama_index.core import QueryBundle
+from llama_index.core.vector_stores import MetadataFilters
 
 logger = logging.getLogger(__name__)
 

diff --git a/executor.py b/executor.py
@@ -9,22 +9,23 @@
 from pathlib import Path
 from urllib.parse import urlparse
 
-from llama_index import ServiceContext, StorageContext
-from llama_index import set_global_service_context
-from llama_index import VectorStoreIndex, SimpleDirectoryReader, Document
-from llama_index.llms import OpenAI
-from llama_index.readers.file.flat_reader import FlatReader
-from llama_index.vector_stores import MilvusVectorStore
-from llama_index.embeddings import HuggingFaceEmbedding
-from llama_index.node_parser.text import SentenceWindowNodeParser
-
-from llama_index.prompts import ChatPromptTemplate, ChatMessage, MessageRole, PromptTemplate
-from llama_index.postprocessor import MetadataReplacementPostProcessor
-from llama_index.postprocessor import SentenceTransformerRerank
-#from llama_index.indices import ZillizCloudPipelineIndex
+from llama_index.core import ServiceContext, StorageContext
+from llama_index.core import set_global_service_context
+from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Document
+from llama_index.llms.openai import OpenAI
+from llama_index.readers.file import FlatReader
+from llama_index.vector_stores.milvus import MilvusVectorStore
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.core.node_parser import SentenceWindowNodeParser
+
+from llama_index.core import ChatPromptTemplate, PromptTemplate
+from llama_index.core.llms import ChatMessage, MessageRole
+from llama_index.core.postprocessor import MetadataReplacementPostProcessor
+from llama_index.core.postprocessor import SentenceTransformerRerank
+from llama_index.indices.managed.zilliz import ZillizCloudPipelineIndex
 from custom.zilliz.base import ZillizCloudPipelineIndex
-from llama_index.indices.query.schema import QueryBundle
-from llama_index.schema import BaseNode, ImageNode, MetadataMode
+from llama_index.core import QueryBundle
+from llama_index.core.schema import BaseNode, ImageNode, MetadataMode
 
 from custom.history_sentence_window import HistorySentenceWindowNodeParser
 from custom.llms.QwenLLM import QwenUnofficial

diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,7 @@
-llama-index==0.9.39
+llama-index==0.10.12
+llama-index-vector-stores-milvus
+llama-index-embeddings-huggingface
+llama-index-indices-managed-zilliz
 pymilvus
 easydict
 dashscope