更新LlamaIndex版本库

2024-09-10 14:07:52 +08:00
parent 6f5548ee61
commit a165d55822
4 changed files with 80 additions and 7 deletions
@@ -0,0 +1,72 @@
+
+from llama_index.llms.xinference import Xinference
+from typing import Any, Callable, Dict, Optional, Sequence, Tuple
+from llama_index.core.llms.callbacks import (
+    llm_chat_callback,
+    llm_completion_callback,
+)
+from llama_index.core.base.llms.types import (
+    ChatMessage,
+    ChatResponse,
+    ChatResponseGen,
+    CompletionResponse,
+    CompletionResponseGen,
+    LLMMetadata,
+    MessageRole,
+)
+from llama_index.llms.xinference.utils import (
+    xinference_message_to_history,
+    xinference_modelname_to_contextsize,
+)
+
+class XinfengModel(Xinference):
+    @llm_chat_callback()
+    def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
+        assert self._generator is not None
+        response_text = self._generator.chat(
+            messages=messages,
+            generate_config={
+                "stream": False,
+                "temperature": self.temperature,
+                "max_tokens": self.max_tokens,
+            },
+        )["choices"][0]["message"]["content"]
+        return ChatResponse(
+            message=ChatMessage(
+                role=MessageRole.ASSISTANT,
+                content=response_text,
+            ),
+            delta=None,
+        )
+
+    @llm_chat_callback()
+    def stream_chat(
+        self, messages: Sequence[ChatMessage], **kwargs: Any
+    ) -> ChatResponseGen:
+        msgs = []
+        for message in messages:
+            msgs.append(message.dict())
+        assert self._generator is not None
+        response_iter = self._generator.chat(
+            messages=msgs,
+            generate_config={
+                "stream": True,
+                "temperature": self.temperature,
+                "max_tokens": self.max_tokens,
+            },
+        )
+
+        def gen() -> ChatResponseGen:
+            text = ""
+            for c in response_iter:
+                delta = c["choices"][0]["delta"].get("content", "")
+                text += delta
+                yield ChatResponse(
+                    message=ChatMessage(
+                        role=MessageRole.ASSISTANT,
+                        content=text,
+                    ),
+                    delta=delta,
+                )
+
+        return gen()