Merge pull request #28 from ericmjl/api-server

API Server

Author: ericmjl

llamabot/bot/chatbot.py

@@ -65,6 +65,8 @@ def __call__(self, message: str) -> AIMessage:
             query=human_message, character_budget=self.response_budget
         )
         messages = [self.system_prompt] + history + [human_message]
+        if self.stream:
+            return self.stream_response(messages)
         response = self.generate_response(messages)
         autorecord(message, response.content)
 

llamabot/bot/imagebot.py

@@ -5,7 +5,7 @@
 import requests
 from pathlib import Path
 from typing import Optional, Union
-from langchain.schema import AIMessage
+from llamabot.components.messages import AIMessage
 
 
 class ImageBot:
@@ -34,6 +34,7 @@ def __call__(
             If it is empty, then we will generate a filename from the prompt.
         :return: The URL of the generated image if running in a Jupyter notebook (str),
             otherwise a pathlib.Path object pointing to the generated image.
+        :raises Exception: If no image URL is found in the response.
         """
         response = self.client.images.generate(
             model=self.model,
@@ -43,6 +44,8 @@ def __call__(
             n=self.n,
         )
         image_url = response.data[0].url
+        if not image_url:
+            raise Exception("No image URL found in response! Please try again.")
 
         # Check if running in a Jupyter notebook
         if is_running_in_jupyter():

llamabot/bot/model_dispatcher.py

@@ -9,7 +9,7 @@
 from llamabot.bot.simplebot import SimpleBot
 from llamabot.components.messages import AIMessage, HumanMessage
 from llamabot.components.docstore import DocumentStore
-
+from llamabot.components.api import APIMixin
 from llamabot.components.messages import (
     RetrievedMessage,
     retrieve_messages_up_to_budget,
@@ -24,7 +24,7 @@
 prompt_recorder_var = contextvars.ContextVar("prompt_recorder")
 
 
-class QueryBot(SimpleBot, DocumentStore):
+class QueryBot(SimpleBot, DocumentStore, APIMixin):
     """QueryBot is a bot that uses simple RAG to answer questions about a document."""
 
     def __init__(

llamabot/bot/querybot.py

@@ -1,6 +1,6 @@
 """Class definition for SimpleBot."""
 import contextvars
-from typing import Optional
+from typing import Optional, Union
 
 
 from llamabot.components.messages import (
@@ -29,6 +29,8 @@ class SimpleBot:
     :param model_name: The name of the model to use.
     :param stream: Whether to stream the output to stdout.
     :param json_mode: Whether to print debug messages.
+    :param api_key: The OpenAI API key to use.
+    :param mock_response: A mock response to use, for testing purposes only.
     """
 
     def __init__(
@@ -39,56 +41,81 @@ def __init__(
         stream=True,
         json_mode: bool = False,
         api_key: Optional[str] = None,
+        mock_response: Optional[str] = None,
     ):
         self.system_prompt: SystemMessage = SystemMessage(content=system_prompt)
         self.model_name = model_name
         self.temperature = temperature
         self.stream = stream
         self.json_mode = json_mode
         self.api_key = api_key
+        self.mock_response = mock_response
 
-    def __call__(self, human_message: str) -> AIMessage:
+    def __call__(self, human_message: str) -> Union[AIMessage, str]:
         """Call the SimpleBot.
 
         :param human_message: The human message to use.
         :return: The response to the human message, primed by the system prompt.
         """
 
-        messages: list[BaseMessage] = [
-            self.system_prompt,
-            HumanMessage(content=human_message),
-        ]
+        messages = [self.system_prompt, HumanMessage(content=human_message)]
+        if self.stream:
+            return self.stream_response(messages)
         response = self.generate_response(messages)
         autorecord(human_message, response.content)
         return response
 
     def generate_response(self, messages: list[BaseMessage]) -> AIMessage:
-        """Generate a response from the given messages."""
-
-        messages_dumped: list[dict] = [m.model_dump() for m in messages]
-        completion_kwargs = dict(
-            model=self.model_name,
-            messages=messages_dumped,
-            temperature=self.temperature,
-            stream=self.stream,
-        )
-        if self.json_mode:
-            completion_kwargs["response_format"] = {"type": "json_object"}
-        if self.api_key:
-            completion_kwargs["api_key"] = self.api_key
-        response = completion(**completion_kwargs)
+        """Generate a response from the given messages.
 
-        if self.stream:
-            ai_message = ""
-            for chunk in response:
-                delta = chunk.choices[0].delta.content
-                if delta is not None:
-                    print(delta, end="")
-                    ai_message += delta
-            return AIMessage(content=ai_message)
+        :param messages: A list of messages.
+        :return: The response to the messages.
+        """
 
+        response = _make_response(self, messages)
         return AIMessage(content=response.choices[0].message.content)
 
+    def stream_response(self, messages: list[BaseMessage]) -> str:
+        """Stream the response from the given messages.
+
+        This is intended to be used with Panel's ChatInterface as part of the callback.
+
+        :param messages: A list of messages.
+        :return: A generator that yields the response.
+        """
+        response = _make_response(self, messages)
+        message = ""
+        for chunk in response:
+            delta = chunk.choices[0].delta.content
+            if delta is not None:
+                message += delta
+                print(delta, end="")
+                yield message
+        print()
+
+
+def _make_response(bot: SimpleBot, messages: list[BaseMessage]):
+    """Make a response from the given messages.
+
+    :param bot: A SimpleBot
+    :param messages: A list of Messages.
+    :return: A response object.
+    """
+    messages_dumped: list[dict] = [m.model_dump() for m in messages]
+    completion_kwargs = dict(
+        model=bot.model_name,
+        messages=messages_dumped,
+        temperature=bot.temperature,
+        stream=bot.stream,
+    )
+    if bot.mock_response:
+        completion_kwargs["mock_response"] = bot.mock_response
+    if bot.json_mode:
+        completion_kwargs["response_format"] = {"type": "json_object"}
+    if bot.api_key:
+        completion_kwargs["api_key"] = bot.api_key
+    return completion(**completion_kwargs)
+
     # Commented out until later.
     # def panel(
     #     self,

llamabot/bot/simplebot.py

@@ -8,7 +8,7 @@
 
 from llamabot import ChatBot, PromptRecorder
 
-from . import blog, configure, doc, git, python, tutorial, zotero, repo
+from . import blog, configure, doc, git, python, tutorial, zotero, repo, serve
 from .utils import exit_if_asked, uniform_prompt
 
 app = typer.Typer()
@@ -39,6 +39,9 @@
 )
 app.add_typer(configure.app, name="configure", help="Configure LlamaBot.")
 app.add_typer(repo.app, name="repo", help="Chat with a code repository.")
+app.add_typer(
+    serve.cli, name="serve", help="Serve up a LlamaBot as a FastAPI endpoint."
+)
 
 
 @app.command()

llamabot/cli/__init__.py

@@ -106,7 +106,8 @@ def compose_commit():
     try:
         diff = get_git_diff()
         bot = commitbot()
-        bot(write_commit_message(diff))
+        msg = bot(write_commit_message(diff))
+        echo(msg.content)
     except Exception as e:
         echo(f"Error encountered: {e}", err=True)
         echo("Please write your own commit message.", err=True)