Basic implementation of an plugin system for OA (#2765)

# the plugins
Hi, this is my first PR here, but I was somewhat active on other fronts
of OA development.

This pr will bring some basic plugin functionality to the Open
Assistant, and as discussed with @yk @andreaskoepf
there are quite a few directions for something like this to be
integrated with OA, but this should serve a purpose as some initial
proof-of-concept and exploratory feature for 3rd party integrations with
OA.

I also included a small **calculator** plugin as a possible candidate
for the OA internal plugin, which would be like the default one, for
people to try out and also as an example, of how one could implement own
plugins.

If we decide to include this plugin, there should be added a
deployment/hosting mechanism for it.

I will push a separate branch in the next couple of days, that will
serve as an alternative to the approach, so we can A/B test it along
with the new models (SFT-s/RLHF-s)

I also tried to comment on every weird quirk or decision in code, so one
could easily understand, and change it, but there are quite a few
places, where a simple new line char or like " char in specific strings,
could affect LLM performance in the plugin usage scenario.

Will also try to push some documentation regarding plugin development,
but there are already some useful comments in **calculator** plugin on
what should be paid attention to.

Here are some of the current UI changes introduced with this PR.
<details>
  <summary>Plugin chooser component</summary>
<img width="854" alt="Screenshot 2023-04-20 at 00 55 38"
src="https://user-images.githubusercontent.com/13547364/233217078-d2e4e28f-36eb-451e-a655-1679188aed52.png">
</details>
<details>
  <summary>Plugin execution details component</summary>
<img width="824" alt="Screenshot 2023-04-19 at 21 40 38"
src="https://user-images.githubusercontent.com/13547364/233216884-e69bcf9c-707f-43de-a52d-41db5d92c504.png">
<img width="744" alt="Screenshot 2023-04-19 at 21 40 56"
src="https://user-images.githubusercontent.com/13547364/233217161-c114f5b9-881f-4476-a2b1-459179a9353e.png">
<img width="545" alt="Screenshot 2023-04-19 at 21 30 18"
src="https://user-images.githubusercontent.com/13547364/233217187-17fb87e5-e4be-43e4-96ac-7cdd84223147.png">
</details>
<details open>
 <summary>Plugin assisted answer</summary>
<img width="837" alt="Screenshot 2023-04-19 at 21 29 52"
src="https://user-images.githubusercontent.com/13547364/233217260-4986f456-efa5-47a5-aabc-926a8a5a9a2f.png">
<img width="943" alt="Screenshot 2023-04-21 at 18 28 45"
src="https://user-images.githubusercontent.com/13547364/233687877-0d0f9ffb-b16a-48de-96ad-e4c3a02f4c66.png">
</details>
<details>
 <summary>Verified plugin usage UI look</summary>
<img width="864" alt="Screenshot 2023-04-20 at 15 08 36"
src="https://user-images.githubusercontent.com/13547364/233376402-52ed5a3d-631a-4350-9130-61548a8d7b02.png">
</details>
<details>
 <summary>Some plugin usage examples</summary>
<img width="1048" alt="Screenshot 2023-04-18 at 01 57 33"
src="https://user-images.githubusercontent.com/13547364/233217685-79b262bd-81fd-4641-9ad9-110e8b689e42.png">
<img width="993" alt="Screenshot 2023-04-17 at 23 17 35"
src="https://user-images.githubusercontent.com/13547364/233217687-561773a1-b16a-49f5-bdbc-f30b46bed33d.png">
</details>
<details open>
<summary>Mixed usage example where model chooses not to use plugin on
its own</summary>
<img width="690" alt="Screenshot 2023-04-20 at 21 31 46"
src="https://user-images.githubusercontent.com/13547364/233469420-25c72893-7c7a-426c-9f4b-ce9144d643ae.png">
</details>

---------

Co-authored-by: agi <you@example.com>
Co-authored-by: Yannic Kilcher <yk@users.noreply.github.com>
Co-authored-by: Oliver Stanley <olivergestanley@gmail.com>

Author: 4 people

inference/full-dev-setup.sh

@@ -13,10 +13,12 @@ else
     INFERENCE_TAG=latest
 fi
 
+POSTGRES_PORT=${POSTGRES_PORT:-5432}
+
 # Creates a tmux window with splits for the individual services
 
 tmux new-session -d -s "inference-dev-setup"
-tmux send-keys "docker run --rm -it -p 5732:5432 -e POSTGRES_PASSWORD=postgres --name postgres postgres" C-m
+tmux send-keys "docker run --rm -it -p $POSTGRES_PORT:5432 -e POSTGRES_PASSWORD=postgres --name postgres postgres" C-m
 tmux split-window -h
 tmux send-keys "docker run --rm -it -p 6779:6379 --name redis redis" C-m
 
@@ -30,7 +32,7 @@ fi
 
 tmux split-window -h
 tmux send-keys "cd server" C-m
-tmux send-keys "LOGURU_LEVEL=$LOGLEVEL POSTGRES_PORT=5732 REDIS_PORT=6779 DEBUG_API_KEYS='0000,0001' ALLOW_DEBUG_AUTH=True TRUSTED_CLIENT_KEYS=6969 uvicorn main:app" C-m
+tmux send-keys "LOGURU_LEVEL=$LOGLEVEL POSTGRES_PORT=$POSTGRES_PORT REDIS_PORT=6779 DEBUG_API_KEYS='0000,0001' ALLOW_DEBUG_AUTH=True TRUSTED_CLIENT_KEYS=6969 uvicorn main:app" C-m
 tmux split-window -h
 tmux send-keys "cd text-client" C-m
 tmux send-keys "sleep 5" C-m

inference/server/alembic/versions/2023_05_01_2253-5b4211625a9f_added_used_plugin_to_message.py

@@ -0,0 +1,28 @@
+"""added used plugin to message
+
+Revision ID: 5b4211625a9f
+Revises: ea19bbc743f9
+Create Date: 2023-05-01 22:53:16.297495
+
+"""
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "5b4211625a9f"
+down_revision = "ea19bbc743f9"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("message", sa.Column("used_plugin", postgresql.JSONB(astext_type=sa.Text()), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("message", "used_plugin")
+    # ### end Alembic commands ###

inference/server/oasst_inference_server/chat_repository.py

@@ -88,12 +88,13 @@ async def abort_work(self, message_id: str, reason: str) -> models.DbMessage:
         await self.session.refresh(message)
         return message
 
-    async def complete_work(self, message_id: str, content: str) -> models.DbMessage:
+    async def complete_work(self, message_id: str, content: str, used_plugin: inference.PluginUsed) -> models.DbMessage:
         logger.debug(f"Completing work on message {message_id}")
         message = await self.get_assistant_message_by_id(message_id)
         message.state = inference.MessageState.complete
         message.work_end_at = datetime.datetime.utcnow()
         message.content = content
+        message.used_plugin = used_plugin
         await self.session.commit()
         logger.debug(f"Completed work on message {message_id}")
         await self.session.refresh(message)

inference/server/oasst_inference_server/database.py

@@ -43,6 +43,8 @@ def custom_json_deserializer(s):
             return chat_schema.CreateMessageRequest.parse_obj(d)
         case "WorkRequest":
             return inference.WorkRequest.parse_obj(d)
+        case "PluginUsed":
+            return inference.PluginUsed.parse_obj(d)
         case None:
             return d
         case _:

inference/server/oasst_inference_server/models/chat.py

@@ -28,6 +28,8 @@ class DbMessage(SQLModel, table=True):
     safety_label: str | None = Field(None)
     safety_rots: str | None = Field(None)
 
+    used_plugin: inference.PluginUsed | None = Field(None, sa_column=sa.Column(pg.JSONB))
+
     state: inference.MessageState = Field(inference.MessageState.manual)
     work_parameters: inference.WorkParameters = Field(None, sa_column=sa.Column(pg.JSONB))
     work_begin_at: datetime.datetime | None = Field(None)
@@ -68,6 +70,7 @@ def to_read(self) -> inference.MessageRead:
             safety_level=self.safety_level,
             safety_label=self.safety_label,
             safety_rots=self.safety_rots,
+            used_plugin=self.used_plugin,
         )
 
 

inference/server/oasst_inference_server/routes/chats.py

@@ -140,6 +140,7 @@ async def create_assistant_message(
             work_parameters = inference.WorkParameters(
                 model_config=model_config,
                 sampling_parameters=request.sampling_parameters,
+                plugins=request.plugins,
             )
             assistant_message = await ucr.initiate_assistant_message(
                 parent_id=request.parent_id,

inference/server/oasst_inference_server/routes/configs.py

@@ -1,9 +1,19 @@
+import asyncio
+
+import aiohttp
 import fastapi
 import pydantic
+import yaml
+from aiohttp.client_exceptions import ClientConnectorError, ServerTimeoutError
+from fastapi import HTTPException
+from loguru import logger
 from oasst_inference_server.settings import settings
 from oasst_shared import model_configs
 from oasst_shared.schemas import inference
 
+# NOTE: Populate this with plugins that we will provide out of the box
+OA_PLUGINS = []
+
 router = fastapi.APIRouter(
     prefix="/configs",
     tags=["configs"],
@@ -63,6 +73,16 @@ class ModelConfigInfo(pydantic.BaseModel):
             repetition_penalty=1.2,
         ),
     ),
+    ParameterConfig(
+        name="k50-Plugins",
+        description="Top-k sampling with k=50 and temperature=0.35",
+        sampling_parameters=inference.SamplingParameters(
+            max_new_tokens=1024,
+            temperature=0.35,
+            top_k=50,
+            repetition_penalty=(1 / 0.90),
+        ),
+    ),
     ParameterConfig(
         name="nucleus9",
         description="Nucleus sampling with p=0.9",
@@ -93,6 +113,44 @@ class ModelConfigInfo(pydantic.BaseModel):
 ]
 
 
+async def fetch_plugin(url: str, retries: int = 3, timeout: float = 5.0) -> inference.PluginConfig:
+    async with aiohttp.ClientSession() as session:
+        for attempt in range(retries):
+            try:
+                async with session.get(url, timeout=timeout) as response:
+                    content_type = response.headers.get("Content-Type")
+
+                    if response.status == 200:
+                        if "application/json" in content_type or url.endswith(".json"):
+                            config = await response.json()
+                        elif (
+                            "application/yaml" in content_type
+                            or "application/x-yaml" in content_type
+                            or url.endswith(".yaml")
+                            or url.endswith(".yml")
+                        ):
+                            config = yaml.safe_load(await response.text())
+                        else:
+                            raise HTTPException(
+                                status_code=400,
+                                detail=f"Unsupported content type: {content_type}. Only JSON and YAML are supported.",
+                            )
+
+                        return inference.PluginConfig(**config)
+                    elif response.status == 404:
+                        raise HTTPException(status_code=404, detail="Plugin not found")
+                    else:
+                        raise HTTPException(status_code=response.status, detail="Unexpected status code")
+            except (ClientConnectorError, ServerTimeoutError) as e:
+                if attempt == retries - 1:  # last attempt
+                    raise HTTPException(status_code=500, detail=f"Request failed after {retries} retries: {e}")
+                await asyncio.sleep(2**attempt)  # exponential backoff
+
+            except aiohttp.ClientError as e:
+                raise HTTPException(status_code=500, detail=f"Request failed: {e}")
+    raise HTTPException(status_code=500, detail="Failed to fetch plugin")
+
+
 @router.get("/model_configs")
 async def get_model_configs() -> list[ModelConfigInfo]:
     return [
@@ -103,3 +161,36 @@ async def get_model_configs() -> list[ModelConfigInfo]:
         for model_config_name in model_configs.MODEL_CONFIGS
         if (settings.allowed_model_config_names == "*" or model_config_name in settings.allowed_model_config_names_list)
     ]
+
+
+@router.post("/plugin_config")
+async def get_plugin_config(plugin: inference.PluginEntry) -> inference.PluginEntry:
+    try:
+        plugin_config = await fetch_plugin(plugin.url)
+    except HTTPException as e:
+        logger.warning(f"Failed to fetch plugin config from {plugin.url}: {e.detail}")
+        raise fastapi.HTTPException(status_code=e.status_code, detail=e.detail)
+
+    return inference.PluginEntry(url=plugin.url, enabled=plugin.enabled, plugin_config=plugin_config)
+
+
+@router.get("/builtin_plugins")
+async def get_builtin_plugins() -> list[inference.PluginEntry]:
+    plugins = []
+
+    for plugin in OA_PLUGINS:
+        try:
+            plugin_config = await fetch_plugin(plugin.url)
+        except HTTPException as e:
+            logger.warning(f"Failed to fetch plugin config from {plugin.url}: {e.detail}")
+            continue
+
+        final_plugin: inference.PluginEntry = inference.PluginEntry(
+            url=plugin.url,
+            enabled=plugin.enabled,
+            trusted=plugin.trusted,
+            plugin_config=plugin_config,
+        )
+        plugins.append(final_plugin)
+
+    return plugins

inference/server/oasst_inference_server/routes/workers.py

@@ -350,6 +350,7 @@ async def handle_generated_text_response(
             message = await cr.complete_work(
                 message_id=message_id,
                 content=response.text,
+                used_plugin=response.used_plugin,
             )
             logger.info(f"Completed work for {message_id=}")
         message_packet = inference.InternalFinishedMessageResponse(

inference/server/oasst_inference_server/schemas/chat.py

@@ -14,6 +14,8 @@ class CreateAssistantMessageRequest(pydantic.BaseModel):
     parent_id: str
     model_config_name: str
     sampling_parameters: inference.SamplingParameters = pydantic.Field(default_factory=inference.SamplingParameters)
+    plugins: list[inference.PluginEntry] = pydantic.Field(default_factory=list[inference.PluginEntry])
+    used_plugin: inference.PluginUsed | None = None
 
 
 class PendingResponseEvent(pydantic.BaseModel):

inference/worker/PLUGINS.md

@@ -0,0 +1,72 @@
+# Plugin system for OA
+
+This is a basic implementation of support for external augmentation and
+OpenAI/ChatGPT plugins into the Open-Assistant. In the current state, this is
+more of a proof-of-concept and should be considered to be used behind some
+experimental flag.
+
+## Architecture
+
+There is now some kind of middleware between work.py(worker) and the final
+prompt that is passed to the inference server for generation and streaming. That
+middleware is responsible for checking if there is an enabled plugin in the
+userland/UI and if so, it will take over the job of creating curated pre-prompts
+for plugin usage, as well as generating subsequent calls to LLM(inner
+monologues) in order to generate the final externally **augmented** prompt, that
+will be passed back to the worker and next to the inference, for final LLM
+generation/streaming tokens to the frontend.
+
+## Plugins
+
+Plugins are in essence just pretty wrappers around some kind of API-s and serve
+a purpose to help LLM utilize it more precisely and reliably, so they can be
+quite useful and powerful augmentation tools for Open-Assistant. Two main parts
+of a plugin are the ai-plugin.json file, which is just the main descriptor of a
+plugin, and the second part is OpenAPI specification of the plugin API-s.
+
+Here is OpenAI plugins
+[specification](https://platform.openai.com/docs/plugins/getting-started) that
+is currently partially supported with this system.
+
+For now, only non-authentication-based and only (**GET** request) plugins are
+supported. Some of them are:
+
+- https://www.klarna.com/.well-known/ai-plugin.json
+- https://www.joinmilo.com/.well-known/ai-plugin.json
+
+Adding support for all other request types would be quite tricky with the
+current approach. It would be best to drop current “mansplaining” of the API to
+LLM and just show it complete json/yaml content. But unfortunately for that to
+be reliable and to work as close as current approach we would need larger
+context size and a bit more capable models.
+
+And quite a few of them can be found on this website
+[plugin "store" wellknown.ai](https://www.wellknown.ai/)
+
+One of the ideas of the plugin system is that we can have some internal OA
+plugins, which will be like out-of-the-box plugins, and there could be endless
+third-party community-developed plugins as well.
+
+### Notes regarding the reliability and performance and the limitations of the plugin system
+
+Performance can vary a lot depending on the models and plugins used. Some of
+them work better some worse, but that aspect should improve as we get better and
+better models. One of the biggest limitations at the moment is context size and
+instruction following capabilities. And that is combated with some prompt
+tricks, truncations of the plugin OpenAPI descriptions and dynamically
+including/excluding parts of the prompts in the internal processing of the
+subsequent generations of intermediate texts (inner monologues). More of the
+limitations and possible alternatives are explained in code comments.
+
+The current approach is somewhat hybrid I would say, and relies on the zero-shot
+capabilities of a model. There will be one more branch with the plugin system
+that will be a bit different approach than this one as it will be utilizing
+other smaller embedding transformer models and vector stores, so we can do A/B
+testing of the system alongside new OA model releases.
+
+## Relevant files for the inference side of the plugin system
+
+- chat_chain.py
+- chat*chain_utils.py *(tweaking tools/plugin description string generation can
+  help for some models)\_
+- chat*chain_prompts.py *(tweaking prompts can help also)\_