Merge pull request #50 from pattern-tech/feat/wallet-accounting

Feat/wallet accounting

Author: yasinfakhar

api/src/conversation/routers/playground_conversation_router.py

@@ -1,8 +1,8 @@
 from uuid import UUID
 from enum import Enum
-from typing import List, Dict
 from pydantic import BaseModel
 from sqlalchemy.orm import Session
+from typing import List, Optional, Dict
 from fastapi.responses import StreamingResponse
 from fastapi import APIRouter, Depends, HTTPException, status
 
@@ -47,6 +47,7 @@ class CreateConversationInput(BaseModel):
     """
     name: str
     project_id: UUID
+    conversation_id: Optional[UUID] = None
 
     class Config:
         from_attributes = True
@@ -82,6 +83,13 @@ class MessageInput(BaseModel):
     stream: bool = True
 
 
+class FirstMessage(BaseModel):
+    """
+    Schema for the first message.
+    """
+    message: str
+
+
 @router.post(
     "",
     response_model=ConversationOutput,
@@ -93,7 +101,7 @@ def create_conversation(
     input: CreateConversationInput,
     db: Session = Depends(get_db),
     service: ConversationService = Depends(get_conversation_service),
-    user_id: UUID = Depends(authenticate_user),
+    user_id: UUID = Depends(authenticate_user)
 ):
     """
     Create a new conversation.
@@ -108,7 +116,7 @@ def create_conversation(
     """
     try:
         conversation = service.create_conversation(
-            db, input.name, input.project_id, user_id)
+            db, input.name, input.project_id, user_id, input.conversation_id)
         return global_response(conversation)
     except Exception as e:
         raise HTTPException(
@@ -326,3 +334,28 @@ async def send_message(
     except Exception as e:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST, detail=str(e))
+
+
+@router.post(
+    "/{project_id}/{conversation_id}/title-generation",
+    summary="Auto Title Generation",
+    description="Using LLM to generate title for conversation",
+    response_description="LLM title generated"
+)
+async def generate_title(
+    input: FirstMessage,
+    project_id: UUID,
+    conversation_id: UUID,
+    db: Session = Depends(get_db),
+    conversation_service: ConversationService = Depends(
+        get_conversation_service),
+    user_id: UUID = Depends(authenticate_user),
+):
+    try:
+        title = conversation_service.rename_title(
+            db, conversation_id, user_id, input.message)
+        return title
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)
+        )

api/src/conversation/services/conversation_service.py

@@ -1,13 +1,13 @@
 from uuid import UUID
 from typing import List
-from datetime import timedelta
 from sqlalchemy.orm import Session
 from fastapi import HTTPException, status
 from langchain_core.messages.human import HumanMessage
 
 from src.util.configuration import Config
 from src.agentflow.agents.hub import AgentHub
 from src.db.models import Conversation, QueryUsage
+from src.agentflow.utils.shared_tools import init_llm
 from src.user.services.user_service import UserService
 from src.agent.services.memory_service import MemoryService
 from src.project.services.project_service import ProjectService
@@ -31,7 +31,7 @@ def __init__(self):
         self.user_service = UserService()
 
     def create_conversation(
-        self, db_session: Session, name: str, project_id: UUID, user_id: UUID
+        self, db_session: Session, name: str, project_id: UUID, user_id: UUID, conversation_id: UUID = None
     ) -> Conversation:
         """
         Creates a new conversation.
@@ -41,6 +41,7 @@ def create_conversation(
             name (str): The name of the conversation.
             project_id (UUID): The ID of the project the conversation belongs to.
             user_id (UUID): The ID of the user creating the conversation.
+            conversation_id (UUID): The ID of the conversation (optional).
 
         Returns:
             Conversation: The created conversation instance.
@@ -53,6 +54,11 @@ def create_conversation(
 
         conversation = Conversation(
             name=name, project_id=project_id, user_id=user_id)
+
+        # conversation id is generated in frontend
+        if conversation_id:
+            conversation.id = conversation_id
+
         return self.repository.create(db_session, conversation)
 
     def get_conversation(
@@ -258,9 +264,46 @@ def get_history(
         memory = self.memory_service.get_memory(conversation_id)
 
         history = []
+        generated_id = 0
         for message in memory.messages:
             history.append({
+                "id": generated_id,
                 "role": "human" if isinstance(message, HumanMessage) else "ai",
                 "content": message.content,
             })
+            generated_id += 1
         return history
+
+    def rename_title(self, db_session: Session, conversation_id: UUID, user_id: UUID, message: str) -> str:
+        """
+        Renames a conversation title using the LLM to generate a title.
+
+        Args:
+            db_session (Session): The database session.
+            conversation_id (UUID): The ID of the conversation to rename.
+            user_id (UUID): The ID of the user renaming the conversation.
+            message (str): The message to generate a title for.
+
+        Returns:
+            str: The new title of the conversation.
+        """
+        config = Config.get_config()
+        self.llm = init_llm(service=config["llm"]["provider"],
+                            model_name=config["llm"]["model"],
+                            api_key=config["llm"]["api_key"],
+                            stream=False,
+                            callbacks=None)
+        messages = [
+            (
+                "system",
+                "You should generate a title (maximum 5 words) for this message",
+            ),
+            ("human", f"{message}"),
+        ]
+
+        title = self.llm.invoke(messages)
+
+        self.repository.update(db_session, conversation_id, {
+                               "name": title.content}, user_id)
+
+        return title.content