[Feat] support thread copy

Author: liuooo

app/api/deps.py

@@ -90,10 +90,17 @@ async def get_param_from_request(request: Request):
     return get_param_from_request
 
 
-def verify_token_relation(relation_type: RelationType, name: str):
+def verify_token_relation(relation_type: RelationType, name: str, ignore_none_relation_id: bool = False):
+    """
+    param relation_type: relation type
+    param name: param name
+    param ignore_none_relation_id: if ignore_none_relation_id is set, return where relation_id is None, use for copy thread api
+    """
     async def verify_authorization(
         session=Depends(get_session), token_id=Depends(get_token_id), relation_id=Depends(get_param(name))
     ):
+        if token_id and ignore_none_relation_id:
+            return
         if token_id and relation_id:
             verify = TokenRelationQuery(token_id=token_id, relation_type=relation_type, relation_id=relation_id)
             if TokenRelationService.verify_relation(session=session, verify=verify):

app/models/thread.py

@@ -16,6 +16,8 @@ class ThreadCreate(BaseModel):
     object: str = "thread"
     messages: Optional[list[MessageCreate]]
     # metadata: Optional[dict]
+    thread_id: Optional[str]
+    end_message_id: Optional[str]
 
 
 class ThreadUpdate(BaseModel):

app/providers/auth_provider.py

@@ -61,7 +61,8 @@ def enable(self):
         verify_thread_depends = Depends(verify_token_relation(relation_type=RelationType.Thread, name="thread_id"))
         for route in thread.router.routes:
             if route.name == thread.create_thread.__name__:
-                route.dependencies.append(Depends(verfiy_token))
+                route.dependencies.append(Depends(
+                    verify_token_relation(relation_type=RelationType.Thread, name="thread_id", ignore_none_relation_id=True)))
             else:
                 route.dependencies.append(verify_thread_depends)
 

app/services/message/message.py

@@ -74,3 +74,21 @@ def get_message_file(*, session: Session, thread_id: str, message_id: str, file_
         if msg_file is None:
             raise ResourceNotFoundError(message="Message file not found")
         return msg_file
+
+    @staticmethod
+    def copy_messages(*, session: Session, from_thread_id: str, to_thread_id: str, end_message_id: str):
+        """
+        copy thread messages to another thread
+        """
+        statement = select(Message).where(Message.thread_id == from_thread_id)
+        if end_message_id:
+            statement = statement.where(Message.id <= end_message_id)
+        original_messages = session.exec(statement.order_by(Message.id))
+
+        for original_message in original_messages:
+            new_message = Message(
+                thread_id=to_thread_id,
+                **original_message.model_dump(exclude={"id", "created_at", "updated_at", "thread_id"})
+            )
+            session.add(new_message)
+        session.commit()

app/services/thread/thread.py

@@ -30,6 +30,14 @@ def create_thread(*, session: Session, body: ThreadCreate, token_id=None) -> Thr
                     thread_id=thread_id,
                     body=MessageCreate.from_orm(message),
                 )
+        elif body.thread_id:
+            # copy thread
+            from app.services.message.message import MessageService
+
+            MessageService.copy_messages(session=session,
+                                         from_thread_id=body.thread_id,
+                                         to_thread_id=thread_id,
+                                         end_message_id=body.end_message_id)
         session.refresh(db_thread)
         return db_thread
 

tests/e2e/test_thread_api.py

@@ -0,0 +1,29 @@
+import pytest
+import openai
+
+@pytest.fixture(name="client")
+def client_for_test():
+    """
+    a openai client connected to local server for test
+    """
+    return openai.OpenAI(base_url="http://localhost:8086/api/v1", api_key="ml-xxx")
+
+
+def test_thread_copy(client):
+    """
+    test copy thread
+    """
+    thread = client.beta.threads.create()
+    contents = ["test1", "test2", "test3"]
+    messages = [client.beta.threads.messages.create(thread_id=thread.id, role="user", content=content) for content in contents]
+    for index, message in enumerate(messages):
+        print(index)
+        new_thread = client.beta.threads.create(extra_body={"thread_id": thread.id,
+                                                        "end_message_id": message.id})
+        new_messages = client.beta.threads.messages.list(thread_id=new_thread.id).data
+        assert len(new_messages) == index + 1
+
+        for i in range(index + 1):
+            assert new_messages[i].content[0].text.value == contents[i]
+        client.beta.threads.delete(thread_id=new_thread.id)
+    client.beta.threads.delete(thread_id=thread.id)