Merge branch 'main' into jerry/bump_version_0.6.38

CHANGELOG.md

@@ -6,6 +6,7 @@
 - Added async support for "compact" and "refine" response modes (#6590)
 - [feature]add transformer tokenize functionalities for optimizer (chinese) (#6659)
 - Add simple benchmark for vector store (#6670)
+- Introduce `llama_index.llms` module, with new `LLM` interface, and `OpenAI`, `HuggingFaceLLM`, `LangChainLLM` implementations. 
 
 ### Bug Fixes / Nits
 - Improve metadata/node storage and retrieval for RedisVectorStore (#6678)

docs/examples/vector_stores/RedisIndexDemo.ipynb

@@ -155,6 +155,7 @@
                 "        self,\n",
                 "        index_name: str,\n",
                 "        index_prefix: str = \"llama_index\",\n",
+                "        prefix_ending: str = \"/vector\",\n",
                 "        index_args: Optional[Dict[str, Any]] = None,\n",
                 "        metadata_fields: Optional[List[str]] = None,\n",
                 "        redis_url: str = \"redis://localhost:6379\",\n",
@@ -174,6 +175,11 @@
                 "        Args:\n",
                 "            index_name (str): Name of the index.\n",
                 "            index_prefix (str): Prefix for the index. Defaults to \"llama_index\".\n",
+                "                The actual prefix used by Redis will be\n",
+                "                \"{index_prefix}{prefix_ending}\".\n",
+                "            prefix_ending (str): Prefix ending for the index. Be careful when\n",
+                "                changing this: https://github.com/jerryjliu/llama_index/pull/6665.\n",
+                "                Defaults to \"/vector\".\n",
                 "            index_args (Dict[str, Any]): Arguments for the index. Defaults to None.\n",
                 "            metadata_fields (List[str]): List of metadata fields to store in the index (only supports TAG fields).\n",
                 "            redis_url (str): URL for the redis instance.\n",

llama_index/llms/__init__.py

@@ -0,0 +1,13 @@
+from llama_index.llms.base import ChatMessage, ChatResponse, ChatResponseGen
+from llama_index.llms.huggingface import HuggingFaceLLM
+from llama_index.llms.langchain import LangChainLLM
+from llama_index.llms.openai import OpenAI
+
+__all__ = [
+    "OpenAI",
+    "LangChainLLM",
+    "HuggingFaceLLM",
+    "ChatMessage",
+    "ChatResponse",
+    "ChatResponseGen",
+]

llama_index/llms/base.py

@@ -0,0 +1,106 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Any, Generator, Optional, Sequence
+
+from pydantic import BaseModel, Field
+
+from llama_index.constants import DEFAULT_CONTEXT_WINDOW, DEFAULT_NUM_OUTPUTS
+
+
+class MessageRole(str, Enum):
+    SYSTEM = "system"
+    USER = "user"
+    ASSISTANT = "assistant"
+    FUNCTION = "function"
+
+
+# ===== Generic Model Input - Chat =====
+class ChatMessage(BaseModel):
+    role: MessageRole = MessageRole.USER
+    content: Optional[str] = ""
+    additional_kwargs: dict = Field(default_factory=dict)
+    name: Optional[str] = None
+
+    def __str__(self) -> str:
+        return f"{self.role.value}: {self.content}"
+
+
+# ===== Generic Model Output - Chat =====
+class ChatResponse(BaseModel):
+    message: ChatMessage
+    raw: Optional[dict] = None
+    delta: Optional[str] = None
+
+    def __str__(self) -> str:
+        return str(self.message)
+
+
+ChatResponseGen = Generator[ChatResponse, None, None]
+
+# ===== Generic Model Output - Completion =====
+class CompletionResponse(BaseModel):
+    text: str
+    additional_kwargs: dict = Field(default_factory=dict)
+    raw: Optional[dict] = None
+    delta: Optional[str] = None
+
+    def __str__(self) -> str:
+        return self.text
+
+
+CompletionResponseGen = Generator[CompletionResponse, None, None]
+
+
+class LLMMetadata(BaseModel):
+    """LLM metadata."""
+
+    context_window: int = DEFAULT_CONTEXT_WINDOW
+    num_output: int = DEFAULT_NUM_OUTPUTS
+
+
+class LLM(ABC):
+    @property
+    @abstractmethod
+    def metadata(self) -> LLMMetadata:
+        pass
+
+    @abstractmethod
+    def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
+        pass
+
+    @abstractmethod
+    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+        pass
+
+    @abstractmethod
+    def stream_chat(
+        self, messages: Sequence[ChatMessage], **kwargs: Any
+    ) -> ChatResponseGen:
+        pass
+
+    @abstractmethod
+    def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+        pass
+
+    # ===== Async Endpoints =====
+    @abstractmethod
+    async def achat(
+        self, messages: Sequence[ChatMessage], **kwargs: Any
+    ) -> ChatResponse:
+        pass
+
+    @abstractmethod
+    async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+        pass
+
+    @abstractmethod
+    async def astream_chat(
+        self, messages: Sequence[ChatMessage], **kwargs: Any
+    ) -> ChatResponseGen:
+        pass
+
+    @abstractmethod
+    async def astream_complete(
+        self, prompt: str, **kwargs: Any
+    ) -> CompletionResponseGen:
+        pass

llama_index/llms/custom.py

@@ -0,0 +1,54 @@
+from typing import Any, Sequence
+
+from llama_index.llms.base import (
+    LLM,
+    ChatMessage,
+    ChatResponse,
+    ChatResponseGen,
+    CompletionResponse,
+    CompletionResponseGen,
+)
+from llama_index.llms.generic_utils import (
+    completion_to_chat_decorator,
+    stream_completion_to_chat_decorator,
+)
+
+
+class CustomLLM(LLM):
+    """Simple abstract base class for custom LLMs.
+
+    Subclasses must implement the `__init__`, `complete`,
+        `stream_complete`, and `metadata` methods.
+    """
+
+    def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
+        chat_fn = completion_to_chat_decorator(self.complete)
+        return chat_fn(messages, **kwargs)
+
+    def stream_chat(
+        self, messages: Sequence[ChatMessage], **kwargs: Any
+    ) -> ChatResponseGen:
+        stream_chat_fn = stream_completion_to_chat_decorator(self.stream_complete)
+        return stream_chat_fn(messages, **kwargs)
+
+    async def achat(
+        self,
+        messages: Sequence[ChatMessage],
+        **kwargs: Any,
+    ) -> ChatResponse:
+        return self.chat(messages, **kwargs)
+
+    async def astream_chat(
+        self,
+        messages: Sequence[ChatMessage],
+        **kwargs: Any,
+    ) -> ChatResponseGen:
+        return self.stream_chat(messages, **kwargs)
+
+    async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+        return self.complete(prompt, **kwargs)
+
+    async def astream_complete(
+        self, prompt: str, **kwargs: Any
+    ) -> CompletionResponseGen:
+        return self.stream_complete(prompt, **kwargs)

llama_index/llms/generic_utils.py

@@ -0,0 +1,154 @@
+from typing import Any, Callable, Sequence
+
+from llama_index.llms.base import (
+    ChatMessage,
+    ChatResponse,
+    ChatResponseGen,
+    CompletionResponse,
+    CompletionResponseGen,
+    MessageRole,
+)
+
+
+def messages_to_prompt(messages: Sequence[ChatMessage]) -> str:
+    """Convert messages to a string prompt."""
+    string_messages = []
+    for message in messages:
+        role = message.role
+        content = message.content
+        string_message = f"{role}: {content}"
+
+        addtional_kwargs = message.additional_kwargs
+        if addtional_kwargs:
+            string_message += f"\n{addtional_kwargs}"
+        string_messages.append(string_message)
+
+    string_messages.append(f"{MessageRole.ASSISTANT}: ")
+    return "\n".join(string_messages)
+
+
+def prompt_to_messages(prompt: str) -> Sequence[ChatMessage]:
+    """Convert a string prompt to a sequence of messages."""
+    return [ChatMessage(role=MessageRole.USER, content=prompt)]
+
+
+def completion_response_to_chat_response(
+    completion_response: CompletionResponse,
+) -> ChatResponse:
+    """Convert a completion response to a chat response."""
+    return ChatResponse(
+        message=ChatMessage(
+            role=MessageRole.ASSISTANT,
+            content=completion_response.text,
+            additional_kwargs=completion_response.additional_kwargs,
+        ),
+        raw=completion_response.raw,
+    )
+
+
+def stream_completion_response_to_chat_response(
+    completion_response_gen: CompletionResponseGen,
+) -> ChatResponseGen:
+    """Convert a stream completion response to a stream chat response."""
+
+    def gen() -> ChatResponseGen:
+        for response in completion_response_gen:
+            yield ChatResponse(
+                message=ChatMessage(
+                    role=MessageRole.ASSISTANT,
+                    content=response.text,
+                    additional_kwargs=response.additional_kwargs,
+                ),
+                delta=response.delta,
+                raw=response.raw,
+            )
+
+    return gen()
+
+
+def chat_response_to_completion_response(
+    chat_response: ChatResponse,
+) -> CompletionResponse:
+    """Convert a chat response to a completion response."""
+    return CompletionResponse(
+        text=chat_response.message.content or "",
+        additional_kwargs=chat_response.message.additional_kwargs,
+        raw=chat_response.raw,
+    )
+
+
+def stream_chat_response_to_completion_response(
+    chat_response_gen: ChatResponseGen,
+) -> CompletionResponseGen:
+    """Convert a stream chat response to a completion response."""
+
+    def gen() -> CompletionResponseGen:
+        for response in chat_response_gen:
+            yield CompletionResponse(
+                text=response.message.content or "",
+                additional_kwargs=response.message.additional_kwargs,
+                delta=response.delta,
+                raw=response.raw,
+            )
+
+    return gen()
+
+
+def completion_to_chat_decorator(
+    func: Callable[..., CompletionResponse]
+) -> Callable[..., ChatResponse]:
+    """Convert a completion function to a chat function."""
+
+    def wrapper(messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
+        # normalize input
+        prompt = messages_to_prompt(messages)
+        completion_response = func(prompt, **kwargs)
+        # normalize output
+        return completion_response_to_chat_response(completion_response)
+
+    return wrapper
+
+
+def stream_completion_to_chat_decorator(
+    func: Callable[..., CompletionResponseGen]
+) -> Callable[..., ChatResponseGen]:
+    """Convert a completion function to a chat function."""
+
+    def wrapper(messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponseGen:
+        # normalize input
+        prompt = messages_to_prompt(messages)
+        completion_response = func(prompt, **kwargs)
+        # normalize output
+        return stream_completion_response_to_chat_response(completion_response)
+
+    return wrapper
+
+
+def chat_to_completion_decorator(
+    func: Callable[..., ChatResponse]
+) -> Callable[..., CompletionResponse]:
+    """Convert a chat function to a completion function."""
+
+    def wrapper(prompt: str, **kwargs: Any) -> CompletionResponse:
+        # normalize input
+        messages = prompt_to_messages(prompt)
+        chat_response = func(messages, **kwargs)
+        # normalize output
+        return chat_response_to_completion_response(chat_response)
+
+    return wrapper
+
+
+def stream_chat_to_completion_decorator(
+    func: Callable[..., ChatResponseGen]
+) -> Callable[..., CompletionResponseGen]:
+    """Convert a chat function to a completion function."""
+
+    def wrapper(prompt: str, **kwargs: Any) -> CompletionResponseGen:
+        # normalize input
+        messages = prompt_to_messages(prompt)
+        chat_response = func(messages, **kwargs)
+        # normalize output
+        return stream_chat_response_to_completion_response(chat_response)
+
+    return wrapper