fix imports and test level (#698)

Author: tastelikefeet

ms_agent/rag/llama_index_rag.py

@@ -2,10 +2,8 @@
 import shutil
 from typing import Any, List, Optional
 
-from llama_index.core.base.llms.types import CompletionResponseGen
 from ms_agent.utils import assert_package_exist
 from omegaconf import DictConfig
-from pydantic import Field
 
 from modelscope import snapshot_download
 from ..llm import LLM, Message
@@ -80,7 +78,7 @@ def complete(_self, prompt: str,
                 def stream_complete(_self,
                                     prompt: str,
                                     formatted: bool = False,
-                                    **kwargs: Any) -> CompletionResponseGen:
+                                    **kwargs: Any):
                     for message in self._llm_instance.generate(
                             messages=[Message(role='user', content=prompt)],
                             stream=True,

tests/rag/test_llama_index.py

@@ -5,6 +5,8 @@
 from ms_agent.rag.llama_index_rag import LlamaIndexRAG
 from omegaconf import DictConfig
 
+from modelscope.utils.test_utils import test_level
+
 
 class LlamaIndexRagTest(unittest.TestCase):
 
@@ -62,6 +64,7 @@ async def retrieve(self):
         result2 = await new_rag.retrieve(query, top_k=3)
         return result1, result2
 
+    @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
     def test_retrieve(self):
         result1, result2 = asyncio.run(self.retrieve())
         self.assertEqual(len(result1), len(result2))