WIP draft of generate() outside of chat.answer()

environment-dev-exl.yml

@@ -0,0 +1,43 @@
+name: ragna-dev-exl
+channels:
+  - pytorch
+  - nvidia
+  - conda-forge
+  - defaults
+dependencies:
+  - python =3.11
+  - pip
+  - git-lfs
+  - jupyterlab>=3
+  - pandas
+  - numpy
+  - panel
+  - tokenizers
+  - pytorch=2.3.1
+  - pytorch-cuda=12.1
+  - cuda-nvcc
+  - rich
+  - pip:
+      - python-dotenv
+      - pytest >=6
+      - pytest-mock
+      - pytest-asyncio
+      - pytest-playwright
+      - mypy ==1.10.0
+      - pre-commit
+      - types-aiofiles
+      - sqlalchemy-stubs
+      - setuptools-scm
+      - pip-tools
+      # documentation
+      - mkdocs
+      - mkdocs-material
+      - mkdocstrings[python]
+      - mkdocs-gen-files
+      - material-plausible-plugin
+      - mkdocs-gallery >=0.10
+      - mdx_truly_sane_lists
+      # exl2
+      - ninja
+      - packaging
+      - exllamav2@https://github.com/turboderp/exllamav2/releases/download/v0.1.5/exllamav2-0.1.5+cu121.torch2.3.1-cp311-cp311-linux_x86_64.whl

ragna/assistants/__init__.py

@@ -18,8 +18,10 @@
     "Jurassic2Ultra",
     "LlamafileAssistant",
     "RagnaDemoAssistant",
+    "Exl2Assistant",
 ]
 
+from ._exl2 import Exl2Assistant
 from ._ai21labs import Jurassic2Ultra
 from ._anthropic import ClaudeHaiku, ClaudeOpus, ClaudeSonnet
 from ._cohere import Command, CommandLight

ragna/assistants/_ai21labs.py

@@ -1,4 +1,4 @@
-from typing import AsyncIterator, cast
+from typing import AsyncIterator, Union, cast
 
 from ragna.core import Message, Source
 
@@ -22,13 +22,51 @@ def _make_system_content(self, sources: list[Source]) -> str:
         )
         return instruction + "\n\n".join(source.content for source in sources)
 
-    async def answer(
-        self, messages: list[Message], *, max_new_tokens: int = 256
+    def _render_prompt(self, prompt: Union[str, list[Message]]) -> Union[str, list]:
+        """
+        Ingests ragna messages-list or a single string prompt and converts to assistant-appropriate format.
+
+        Returns:
+            ordered list of dicts with 'text' and 'role' keys
+        """
+        if isinstance(prompt, str):
+            return [
+                {
+                    "text": prompt,
+                    "role": "user",
+                }
+            ]
+        else:
+            messages = [
+                {"text": i["content"], "role": i["role"]}
+                for i in prompt
+                if i["role"] != "system"
+            ]
+            return messages
+
+    async def generate(
+        self,
+        prompt: Union[str, list[Message]],
+        system_prompt: str,
+        *,
+        max_new_tokens: int = 256,
     ) -> AsyncIterator[str]:
+        """
+        Primary method for calling assistant inference, either as a one-off request from anywhere in ragna, or as part of self.answer()
+        This method should be called for tasks like pre-processing, agentic tasks, or any other user-defined calls.
+
+        Args:
+            prompt: Either a single prompt string or a list of ragna messages
+            system_prompt: System prompt string
+            max_new_tokens: Max number of completion tokens (default 256_
+
+        Returns:
+            async streamed inference response string chunks
+        """
         # See https://docs.ai21.com/reference/j2-chat-api#chat-api-parameters
         # See https://docs.ai21.com/reference/j2-complete-api-ref#api-parameters
         # See https://docs.ai21.com/reference/j2-chat-api#understanding-the-response
-        prompt, sources = (message := messages[-1]).content, message.sources
+
         async for data in self._call_api(
             "POST",
             f"https://api.ai21.com/studio/v1/j2-{self._MODEL_TYPE}/chat",
@@ -41,17 +79,21 @@ async def answer(
                 "numResults": 1,
                 "temperature": 0.0,
                 "maxTokens": max_new_tokens,
-                "messages": [
-                    {
-                        "text": prompt,
-                        "role": "user",
-                    }
-                ],
-                "system": self._make_system_content(sources),
+                "messages": _render_prompt(prompt),
+                "system": system_prompt,
             },
         ):
             yield cast(str, data["outputs"][0]["text"])
 
+    async def answer(
+        self, messages: list[Message], *, max_new_tokens: int = 256
+    ) -> AsyncIterator[str]:
+        prompt, sources = (message := messages[-1]).content, message.sources
+        system_prompt = self._make_system_content(sources)
+        yield generate(
+            prompt=prompt, system_prompt=system_prompt, max_new_tokens=max_new_tokens
+        )
+
 
 # The Jurassic2Mid assistant receives a 500 internal service error from the remote
 # server. See https://github.com/Quansight/ragna/pull/303

ragna/assistants/_anthropic.py

@@ -1,4 +1,4 @@
-from typing import AsyncIterator, cast
+from typing import AsyncIterator, Union, cast
 
 from ragna.core import Message, PackageRequirement, RagnaException, Requirement, Source
 
@@ -36,12 +36,50 @@ def _instructize_system_prompt(self, sources: list[Source]) -> str:
             + "</documents>"
         )
 
-    async def answer(
-        self, messages: list[Message], *, max_new_tokens: int = 256
+    def _render_prompt(self, prompt: Union[str, list[Message]]) -> Union[str, list]:
+        """
+        Ingests ragna messages-list or a single string prompt and converts to assistant-appropriate format.
+
+        Returns:
+            ordered list of dicts with 'content' and 'role' keys
+        """
+        if isinstance(prompt, str):
+            return [
+                {
+                    "content": prompt,
+                    "role": "user",
+                }
+            ]
+        else:
+            messages = [
+                {"content": i["content"], "role": i["role"]}
+                for i in prompt
+                if i["role"] != "system"
+            ]
+            return messages
+
+    async def generate(
+        self,
+        prompt: Union[str, list[Message]],
+        system_prompt: str,
+        *,
+        max_new_tokens: int = 256,
     ) -> AsyncIterator[str]:
+        """
+        Primary method for calling assistant inference, either as a one-off request from anywhere in ragna, or as part of self.answer()
+        This method should be called for tasks like pre-processing, agentic tasks, or any other user-defined calls.
+
+        Args:
+            prompt: Either a single prompt string or a list of ragna messages
+            system_prompt: System prompt string
+            max_new_tokens: Max number of completion tokens (default 256)
+
+        Returns:
+            async streamed inference response string chunks
+        """
         # See https://docs.anthropic.com/claude/reference/messages_post
         # See https://docs.anthropic.com/claude/reference/streaming
-        prompt, sources = (message := messages[-1]).content, message.sources
+
         async for data in self._call_api(
             "POST",
             "https://api.anthropic.com/v1/messages",
@@ -53,8 +91,8 @@ async def answer(
             },
             json={
                 "model": self._MODEL,
-                "system": self._instructize_system_prompt(sources),
-                "messages": [{"role": "user", "content": prompt}],
+                "system": system,
+                "messages": _render_prompt(prompt),
                 "max_tokens": max_new_tokens,
                 "temperature": 0.0,
                 "stream": True,
@@ -70,6 +108,15 @@ async def answer(
 
             yield cast(str, data["delta"].pop("text"))
 
+    async def answer(
+        self, messages: list[Message], *, max_new_tokens: int = 256
+    ) -> AsyncIterator[str]:
+        prompt, sources = (message := messages[-1]).content, message.sources
+        system_prompt = self._instructize_system_prompt(sources)
+        yield self.generate(
+            prompt=prompt, system_prompt=system_prompt, max_new_tokens=max_new_tokens
+        )
+
 
 class ClaudeOpus(AnthropicAssistant):
     """[Claude 3 Opus](https://docs.anthropic.com/claude/docs/models-overview)

ragna/assistants/_cohere.py

@@ -1,4 +1,4 @@
-from typing import AsyncIterator, cast
+from typing import AsyncIterator, Union, cast
 
 from ragna.core import Message, RagnaException, Source
 
@@ -24,13 +24,44 @@ def _make_preamble(self) -> str:
     def _make_source_documents(self, sources: list[Source]) -> list[dict[str, str]]:
         return [{"title": source.id, "snippet": source.content} for source in sources]
 
-    async def answer(
-        self, messages: list[Message], *, max_new_tokens: int = 256
+    def _render_prompt(self, prompt: Union[str, list[Message]]) -> str:
+        """
+        Ingests ragna messages-list or a single string prompt and converts to assistant-appropriate format.
+
+        Returns:
+            prompt string
+        """
+        if isinstance(prompt, str):
+            return prompt
+        else:
+            messages = [i["content"] for i in prompt if i["role"] == "user"][-1]
+            return messages
+
+    async def generate(
+        self,
+        prompt: Union[str, list[Message]],
+        system_prompt: str,
+        source_documents: list[dict[str, str]],
+        *,
+        max_new_tokens: int = 256,
     ) -> AsyncIterator[str]:
+        """
+        Primary method for calling assistant inference, either as a one-off request from anywhere in ragna, or as part of self.answer()
+        This method should be called for tasks like pre-processing, agentic tasks, or any other user-defined calls.
+
+        Args:
+            prompt: Either a single prompt string or a list of ragna messages
+            system_prompt: System prompt string
+            source_documents: List of source content dicts with 'title' and 'snippet' keys
+            max_new_tokens: Max number of completion tokens (default 256)
+
+        Returns:
+            async streamed inference response string chunks
+        """
         # See https://docs.cohere.com/docs/cochat-beta
         # See https://docs.cohere.com/reference/chat
         # See https://docs.cohere.com/docs/retrieval-augmented-generation-rag
-        prompt, sources = (message := messages[-1]).content, message.sources
+
         async for event in self._call_api(
             "POST",
             "https://api.cohere.ai/v1/chat",
@@ -40,8 +71,8 @@ async def answer(
                 "authorization": f"Bearer {self._api_key}",
             },
             json={
-                "preamble_override": self._make_preamble(),
-                "message": prompt,
+                "preamble_override": system_prompt,
+                "message": _render_prompt(prompt),
                 "model": self._MODEL,
                 "stream": True,
                 "temperature": 0.0,
@@ -57,6 +88,22 @@ async def answer(
             if "text" in event:
                 yield cast(str, event["text"])
 
+    async def answer(
+        self, messages: list[Message], *, max_new_tokens: int = 256
+    ) -> AsyncIterator[str]:
+        # See https://docs.cohere.com/docs/cochat-beta
+        # See https://docs.cohere.com/reference/chat
+        # See https://docs.cohere.com/docs/retrieval-augmented-generation-rag
+        prompt, sources = (message := messages[-1]).content, message.sources
+        system_prompt = self._make_preamble()
+        source_documents = self._make_source_documents(sources)
+        yield generate(
+            prompt=prompt,
+            system_prompt=system_prompt,
+            source_documents=source_documents,
+            max_new_tokens=max_new_tokens,
+        )
+
 
 class Command(CohereAssistant):
     """