Support for custom agent in conv

.env.template

@@ -26,4 +26,5 @@ EMAIL_FROM_ADDRESS=
 RESEND_API_KEY=
 ANTHROPIC_API_KEY=
 POSTHOG_API_KEY=
-POSTHOG_HOST=
+POSTHOG_HOST=
+POTPIE_PLUS_BASE_URL=http://localhost:8080

GETTING_STARTED.md

@@ -113,7 +113,7 @@ curl -X 'POST' \
 ```
 ## Step 6: Send Messages in a Conversation
 
-This API returns a stream response for the 
+This API returns a stream response for the
 ```bash
 curl -X 'POST' \
   'http://localhost:8001/api/v1/conversations/1234/message/' \

app/alembic/versions/20241028204107_684a330f9e9f_new_migration.py

@@ -5,30 +5,41 @@
 Create Date: 2024-10-28 20:41:07.469748
 
 """
+
 from typing import Sequence, Union
 
-from alembic import op
 import sqlalchemy as sa
-from sqlalchemy.dialects import postgresql
+from alembic import op
 
 # revision identifiers, used by Alembic.
-revision: str = '20241028204107_684a330f9e9f'
-down_revision: Union[str, None] = '20241003153813_827623103002'
+revision: str = "20241028204107_684a330f9e9f"
+down_revision: Union[str, None] = "20241003153813_827623103002"
 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.execute("CREATE TYPE visibility AS ENUM ('PRIVATE', 'PUBLIC')")
-    op.add_column('conversations', sa.Column('visibility', sa.Enum('PRIVATE', 'PUBLIC', name='visibility'), nullable=False, server_default='PRIVATE'))
+    op.add_column(
+        "conversations",
+        sa.Column(
+            "visibility",
+            sa.Enum("PRIVATE", "PUBLIC", name="visibility"),
+            nullable=False,
+            server_default="PRIVATE",
+        ),
+    )
 
     # ### end Alembic commands ###
 
 
 def downgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    op.add_column('conversations', sa.Column('is_public', sa.BOOLEAN(), autoincrement=False, nullable=True))
-    op.drop_column('conversations', 'visibility')
+    op.add_column(
+        "conversations",
+        sa.Column("is_public", sa.BOOLEAN(), autoincrement=False, nullable=True),
+    )
+    op.drop_column("conversations", "visibility")
     op.execute("DROP TYPE visibility")
     # ### end Alembic commands ###

app/main.py

@@ -19,6 +19,7 @@
 from app.modules.intelligence.prompts.prompt_router import router as prompt_router
 from app.modules.intelligence.prompts.system_prompt_setup import SystemPromptSetup
 from app.modules.intelligence.provider.provider_router import router as provider_router
+from app.modules.intelligence.tools.tool_router import router as tool_router
 from app.modules.key_management.secret_manager import router as secret_manager_router
 from app.modules.parsing.graph_construction.parsing_router import (
     router as parsing_router,
@@ -96,6 +97,7 @@ def include_routers(self):
         self.app.include_router(agent_router, prefix="/api/v1", tags=["Agents"])
 
         self.app.include_router(provider_router, prefix="/api/v1", tags=["Providers"])
+        self.app.include_router(tool_router, prefix="/api/v1", tags=["Tools"])
 
     def add_health_check(self):
         @self.app.get("/health", tags=["Health"])

app/modules/auth/auth_router.py

@@ -10,11 +10,11 @@
 
 from app.core.database import get_db
 from app.modules.auth.auth_schema import LoginRequest
+from app.modules.auth.auth_service import auth_handler
 from app.modules.users.user_schema import CreateUser
 from app.modules.users.user_service import UserService
 from app.modules.utils.APIRouter import APIRouter
 from app.modules.utils.posthog_helper import PostHogClient
-from app.modules.auth.auth_service import auth_handler
 
 SLACK_WEBHOOK_URL = os.getenv("SLACK_WEBHOOK_URL", None)
 
@@ -29,7 +29,6 @@ async def send_slack_message(message: str):
 
 
 class AuthAPI:
-
     @auth_router.post("/login")
     async def login(login_request: LoginRequest):
         email, password = login_request.email, login_request.password
@@ -39,10 +38,8 @@ async def login(login_request: LoginRequest):
             id_token = res.get("idToken")
             return JSONResponse(content={"token": id_token}, status_code=200)
         except Exception as e:
-            return JSONResponse(
-                content={"error": f"ERROR: {str(e)}"}, status_code=400
-            )
-
+            return JSONResponse(content={"error": f"ERROR: {str(e)}"}, status_code=400)
+
     @auth_router.post("/signup")
     async def signup(request: Request, db: Session = Depends(get_db)):
         body = json.loads(await request.body())

app/modules/auth/auth_service.py

@@ -1,3 +1,6 @@
+import hashlib
+import hmac
+import json
 import logging
 import os
 
@@ -65,5 +68,31 @@ async def check_auth(
             res.headers["WWW-Authenticate"] = 'Bearer realm="auth_required"'
             return decoded_token
 
+    @staticmethod
+    def verify_hmac_signature(payload_body: dict, hmac_signature: str) -> bool:
+        if os.getenv("ENV") == "development":
+            return True
+        else:
+            shared_key = AuthService.get_hmac_secret_key()
+            expected_signature = hmac.new(
+                key=shared_key.encode(),
+                msg=json.dumps(payload_body).encode(),
+                digestmod=hashlib.sha256,
+            ).hexdigest()
+            return hmac.compare_digest(hmac_signature, expected_signature)
+
+    @staticmethod
+    def verify_hmac_signature_for_get(user_id: str, hmac_signature: str) -> bool:
+        secret_key = AuthService.get_hmac_secret_key()
+        message = f"user_id={user_id}"
+        expected_signature = hmac.new(
+            secret_key.encode(), message.encode(), hashlib.sha256
+        ).hexdigest()
+        return hmac.compare_digest(expected_signature, hmac_signature)
+
+    @staticmethod
+    def get_hmac_secret_key() -> str:
+        return os.environ.get("SHARED_HMAC_KEY")
+
 
 auth_handler = AuthService()

app/modules/conversations/access/access_schema.py

@@ -1,14 +1,16 @@
-from typing import List
+from typing import List, Optional
 
 from pydantic import BaseModel, EmailStr
-from typing import Optional, List
+
 from app.modules.conversations.conversation.conversation_model import Visibility
 
+
 class ShareChatRequest(BaseModel):
-    conversation_id: str  
-    recipientEmails: Optional[List[EmailStr]]=None
+    conversation_id: str
+    recipientEmails: Optional[List[EmailStr]] = None
     visibility: Visibility
 
+
 class ShareChatResponse(BaseModel):
     message: str
     sharedID: str
@@ -17,5 +19,6 @@ class ShareChatResponse(BaseModel):
 class SharedChatResponse(BaseModel):
     chat: dict
 
+
 class RemoveAccessRequest(BaseModel):
     emails: List[EmailStr]

app/modules/conversations/access/access_service.py

@@ -1,12 +1,13 @@
 from typing import List
-from uuid6 import uuid7 
-from app.modules.conversations.conversation.conversation_model import Conversation, Visibility
-from fastapi import HTTPException
 
+from fastapi import HTTPException
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session
 
-from app.modules.conversations.conversation.conversation_model import Conversation
+from app.modules.conversations.conversation.conversation_model import (
+    Conversation,
+    Visibility,
+)
 
 
 class ShareChatServiceError(Exception):
@@ -17,16 +18,28 @@ class ShareChatService:
     def __init__(self, db: Session):
         self.db = db
 
-    async def share_chat(self, conversation_id: str, user_id: str, recipient_emails: List[str] = None, visibility: Visibility = Visibility.PRIVATE) -> str:
-        chat = self.db.query(Conversation).filter_by(id=conversation_id, user_id=user_id).first()
+    async def share_chat(
+        self,
+        conversation_id: str,
+        user_id: str,
+        recipient_emails: List[str] = None,
+        visibility: Visibility = Visibility.PRIVATE,
+    ) -> str:
+        chat = (
+            self.db.query(Conversation)
+            .filter_by(id=conversation_id, user_id=user_id)
+            .first()
+        )
         if not chat:
-            raise HTTPException(404,"Chat does not exist or you are not authorized to access it.")
-
+            raise HTTPException(
+                404, "Chat does not exist or you are not authorized to access it."
+            )
+
         if visibility == Visibility.PUBLIC:
             chat.visibility = Visibility.PUBLIC
             self.db.commit()
             return conversation_id
-        
+
         if visibility == Visibility.PRIVATE:
             chat.visibility = Visibility.PRIVATE
             if recipient_emails:
@@ -35,66 +48,89 @@ async def share_chat(self, conversation_id: str, user_id: str, recipient_emails:
                 unique_new_emails_set = set(recipient_emails)
 
                 if unique_new_emails_set.issubset(existing_emails_set):
-                    raise ShareChatServiceError("All provided emails have already been shared.")
+                    raise ShareChatServiceError(
+                        "All provided emails have already been shared."
+                    )
 
                 to_share = unique_new_emails_set - existing_emails_set
                 if to_share:
                     try:
                         updated_emails = existing_emails + list(to_share)
-                        self.db.query(Conversation).filter_by(id=conversation_id).update(
-                        {Conversation.shared_with_emails: updated_emails, Conversation.visibility: visibility},
-                        synchronize_session=False  
+                        self.db.query(Conversation).filter_by(
+                            id=conversation_id
+                        ).update(
+                            {
+                                Conversation.shared_with_emails: updated_emails,
+                                Conversation.visibility: visibility,
+                            },
+                            synchronize_session=False,
                         )
-                        self.db.commit()        
+                        self.db.commit()
                     except IntegrityError as e:
                         self.db.rollback()
-                        raise ShareChatServiceError("Failed to update shared chat due to a database integrity error.") from e
+                        raise ShareChatServiceError(
+                            "Failed to update shared chat due to a database integrity error."
+                        ) from e
                 self.db.commit()
                 return conversation_id
             else:
-                    self.db.query(Conversation).filter_by(id=conversation_id).update(
-                        {Conversation.visibility: visibility},
-                        synchronize_session=False  
-                        )
-                    self.db.commit()  
+                self.db.query(Conversation).filter_by(id=conversation_id).update(
+                    {Conversation.visibility: visibility}, synchronize_session=False
+                )
+                self.db.commit()
 
         return conversation_id
 
     async def get_shared_emails(self, conversation_id: str, user_id: str) -> List[str]:
-
-        chat = self.db.query(Conversation).filter_by(id=conversation_id, user_id=user_id).first()
+        chat = (
+            self.db.query(Conversation)
+            .filter_by(id=conversation_id, user_id=user_id)
+            .first()
+        )
         if not chat:
-            raise HTTPException(404,"Chat does not exist or you are not authorized to access it.")
-
+            raise HTTPException(
+                404, "Chat does not exist or you are not authorized to access it."
+            )
+
         return chat.shared_with_emails or []
 
-    async def remove_access(self, conversation_id: str, user_id: str, emails_to_remove: List[str]) -> bool:
+    async def remove_access(
+        self, conversation_id: str, user_id: str, emails_to_remove: List[str]
+    ) -> bool:
         """Remove access for specified emails from a conversation."""
-        chat = self.db.query(Conversation).filter_by(id=conversation_id, user_id=user_id).first()
+        chat = (
+            self.db.query(Conversation)
+            .filter_by(id=conversation_id, user_id=user_id)
+            .first()
+        )
         if not chat:
             raise HTTPException(
                 status_code=404,
-                detail="Chat does not exist or you are not authorized to access it."
+                detail="Chat does not exist or you are not authorized to access it.",
             )
-        
+
         if not chat.shared_with_emails:
             raise ShareChatServiceError("Chat has no shared access to remove.")
-            
+
         existing_emails = set(chat.shared_with_emails)
         emails_to_remove_set = set(emails_to_remove)
-        
+
         # Check if any of the emails to remove actually have access
         if not emails_to_remove_set.intersection(existing_emails):
-            raise ShareChatServiceError("None of the specified emails have access to this chat.")
-
+            raise ShareChatServiceError(
+                "None of the specified emails have access to this chat."
+            )
+
         try:
             updated_emails = list(existing_emails - emails_to_remove_set)
             self.db.query(Conversation).filter_by(id=conversation_id).update(
                 {Conversation.shared_with_emails: updated_emails},
-                synchronize_session=False
+                synchronize_session=False,
             )
             self.db.commit()
             return True
         except IntegrityError as e:
             self.db.rollback()
-            raise ShareChatServiceError("Failed to update shared chat due to a database integrity error.") from e
+            raise ShareChatServiceError(
+                "Failed to update shared chat due to a database integrity error."
+            ) from e

app/modules/conversations/conversation/conversation_model.py

@@ -2,7 +2,7 @@
 
 from sqlalchemy import ARRAY, TIMESTAMP, Column
 from sqlalchemy import Enum as SQLAEnum
-from sqlalchemy import ForeignKey, String, func, Boolean
+from sqlalchemy import ForeignKey, String, func
 from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.orm import relationship
 
@@ -15,10 +15,12 @@ class ConversationStatus(enum.Enum):
     ARCHIVED = "archived"
     DELETED = "deleted"
 
+
 class Visibility(enum.Enum):
     PRIVATE = "private"
     PUBLIC = "public"
 
+
 class Conversation(Base):
     __tablename__ = "conversations"
 
@@ -43,7 +45,9 @@ class Conversation(Base):
         nullable=False,
     )
     shared_with_emails = Column(ARRAY(String), nullable=True)
-    visibility = Column(SQLAEnum(Visibility), default=Visibility.PRIVATE, nullable=False)
+    visibility = Column(
+        SQLAEnum(Visibility), default=Visibility.PRIVATE, nullable=False
+    )
     # Relationships
     user = relationship("User", back_populates="conversations")
     messages = relationship(