feat: add extra body to assistant

Author: YangZhiBoGreenHand

app/api/routes.py

@@ -13,4 +13,3 @@ def router_init():
     api_router.include_router(files.router, prefix="/files", tags=["files"])
     api_router.include_router(token.router, prefix="/tokens", tags=["tokens"])
     api_router.include_router(action.router, prefix="/actions", tags=["actions"])
-

app/core/runner/llm_backend.py

@@ -13,22 +13,25 @@ class LLMBackend:
     def __init__(self, base_url: str, api_key) -> None:
         self.base_url = base_url + "/" if base_url else None
         self.api_key = api_key
-        self.client = OpenAI(
-          base_url=self.base_url,
-          api_key=self.api_key 
-        )
+        self.client = OpenAI(base_url=self.base_url, api_key=self.api_key)
 
     def run(
-        self, messages: List, model: str, tools: List = None, tool_choice="auto", stream=False
+        self, messages: List, model: str, tools: List = None, tool_choice="auto", stream=False, extra_body=None
     ) -> ChatCompletion | Stream[ChatCompletionChunk]:
         chat_params = {
             "messages": messages,
             "model": model,
             "stream": stream,
         }
+        if extra_body:
+            model_params = extra_body.get("model_params")
+            if model_params:
+                if "n" in model_params:
+                    raise ValueError("n is not allowed in model_params")
+                chat_params.update(model_params)
         if tools:
-            chat_params['tools'] = tools
-            chat_params['tool_choice'] = tool_choice if tool_choice else "auto"
+            chat_params["tools"] = tools
+            chat_params["tool_choice"] = tool_choice if tool_choice else "auto"
         logging.info("chat_params: %s", chat_params)
         response = self.client.chat.completions.create(**chat_params)
         logging.info("chat_response: %s", response)

app/core/runner/thread_runner.py

@@ -106,6 +106,7 @@ def __run_step(self, llm: LLMBackend, run: Run, run_steps: List[RunStep], instru
             tools=[tool.openai_function for tool in tools],
             tool_choice="auto" if len(run_steps) < self.max_step else "none",
             stream=True,
+            extra_body=run.extra_body,
         )
 
         # create message creation run step callback

app/models/action.py

@@ -1,7 +1,7 @@
 from sqlalchemy import Column, JSON
 from typing import Optional
 
-from sqlmodel import Field, Boolean
+from sqlmodel import Field
 
 from app.models.base_model import BaseModel, TimeStampMixin, PrimaryKeyMixin
 
@@ -27,5 +27,4 @@ class ActionBase(BaseModel):
 
 
 class Action(ActionBase, PrimaryKeyMixin, TimeStampMixin, table=True):
-
     object: str = Field(nullable=False, default="action")

app/models/assistant.py

@@ -14,6 +14,7 @@ class AssistantBase(BaseModel):
     metadata_: Optional[dict] = Field(default=None, sa_column=Column("metadata", JSON))
     name: Optional[str] = Field(default=None)
     tools: Optional[list] = Field(default=None, sa_column=Column(JSON))
+    extra_body: Optional[dict] = Field(default=None, sa_column=Column(JSON))
 
 
 class Assistant(AssistantBase, PrimaryKeyMixin, TimeStampMixin, table=True):

app/models/run.py

@@ -1,4 +1,4 @@
-from typing import Optional, Any
+from typing import Optional, Any, Union
 
 from sqlalchemy import Column, Enum
 from sqlalchemy.sql.sqltypes import JSON, TEXT
@@ -61,7 +61,7 @@ class RunCreate(BaseModel):
     file_ids: Optional[list] = []
     metadata_: Optional[dict] = Field(default={}, alias="metadata")
     tools: Optional[list] = []
-    extra_body: Optional[dict[str, dict[str, Authentication]]] = {}
+    extra_body: Optional[dict[str, Union[dict[str, Union[Authentication, Any]], Any]]] = {}
 
     @root_validator()
     def root_validator(cls, data: Any):

app/services/run/run.py

@@ -30,6 +30,8 @@ def create_run(
             body.instructions = db_asst.instructions
         if not body.tools and db_asst.tools:
             body.tools = db_asst.tools
+        if not body.extra_body and db_asst.extra_body:
+            body.extra_body = db_asst.extra_body
         # create run
         db_run = Run.model_validate(body.model_dump(), update={"thread_id": thread_id, "file_ids": db_asst.file_ids})
         session.add(db_run)

migrations/versions/2024-03-19-15-27_8dbb8f38ef77.py

@@ -0,0 +1,31 @@
+"""add extra body to assistants
+
+Revision ID: 8dbb8f38ef77
+Revises: e7339aab6549
+Create Date: 2024-03-19 15:27:39.793603
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel
+
+
+# revision identifiers, used by Alembic.
+revision: str = '8dbb8f38ef77'
+down_revision: Union[str, None] = 'e7339aab6549'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('assistant', sa.Column('extra_body', sa.JSON(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('assistant', 'extra_body')
+    # ### end Alembic commands ###

tests/e2e/assistant_test.py

@@ -0,0 +1,46 @@
+import openai
+
+from app.api.deps import get_session
+from app.services.assistant.assistant import AssistantService
+
+
+# 测试创建动作
+def test_create_assistant():
+    client = openai.OpenAI(base_url="http://localhost:8086/api/v1", api_key="xxx")
+    assistant = client.beta.assistants.create(
+        name="Assistant Demo",
+        instructions="你是一个有用的助手",
+        extra_body={
+            "extra_body": {
+                "frequency_penalty": 0,
+                "logit_bias": None,
+                "logprobs": True,
+                "top_logprobs": 0,
+                "max_tokens": 1024,
+                "presence_penalty": 0.6,
+                "temperature": 1,
+                "n": 1,
+                "presence_penalty": 0,
+                "top_p": 1,
+            }
+        },
+        # https://platform.openai.com/docs/api-reference/chat/create 具体参数看这里
+        model="gpt-3.5-turbo-1106",
+    )
+    session = next(get_session())
+    assistant = AssistantService.get_assistant(session=session, assistant_id=assistant.id)
+    assert assistant.name == "Assistant Demo"
+    assert assistant.instructions == "你是一个有用的助手"
+    assert assistant.model == "gpt-3.5-turbo-1106"
+    assert assistant.extra_body == {
+        "frequency_penalty": 0,
+        "logit_bias": None,
+        "logprobs": True,
+        "top_logprobs": 0,
+        "max_tokens": 1024,
+        "presence_penalty": 0.6,
+        "temperature": 1,
+        "n": 1,
+        "presence_penalty": 0,
+        "top_p": 1,
+    }

tests/run/run_with_assistant_extra_body_test.py

@@ -0,0 +1,63 @@
+import time
+
+import openai
+
+from app.api.deps import get_session
+from app.schemas.tool.action import ActionBulkCreateRequest
+from app.schemas.tool.authentication import Authentication, AuthenticationType
+from app.services.tool.action import ActionService
+
+
+def test_run_with_assistant_extra_body():
+    client = openai.OpenAI(base_url="http://localhost:8086/api/v1", api_key="xxx")
+    # 创建带有 action 的 assistant
+    assistant = client.beta.assistants.create(
+        name="Assistant Demo",
+        instructions="你是一个有用的助手",
+        model="gpt-3.5-turbo-1106",
+        extra_body={
+            "extra_body": {
+                "model_params": {
+                    "frequency_penalty": 0,
+                    "logit_bias": None,
+                    "max_tokens": 1024,
+                    "presence_penalty": 0.6,
+                    "temperature": 1,
+                    "presence_penalty": 0,
+                    "top_p": 1,
+                }
+            }
+        },
+    )
+    print(assistant, end="\n\n")
+
+    thread = client.beta.threads.create()
+    print(thread, end="\n\n")
+
+    message = client.beta.threads.messages.create(
+        thread_id=thread.id,
+        role="user",
+        content="你好,介绍一下你自己",
+    )
+    print(message, end="\n\n")
+
+    run = client.beta.threads.runs.create(thread_id=thread.id, assistant_id=assistant.id, instructions="")
+    print(run, end="\n\n")
+
+    while True:
+        # run = client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id)
+        run = client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id)
+        if run.status == "completed":
+            print("done!", end="\n\n")
+            messages = client.beta.threads.messages.list(thread_id=thread.id)
+
+            print("messages: ")
+            for message in messages:
+                assert message.content[0].type == "text"
+                print(messages)
+                print({"role": message.role, "message": message.content[0].text.value})
+
+            break
+        else:
+            print("\nin progress...")
+            time.sleep(1)