Add messages limitation

apis/LangSmith.py

@@ -0,0 +1,72 @@
+import requests
+import json
+import os 
+from time import sleep
+from dotenv import load_dotenv
+from datetime import datetime
+
+
+load_dotenv()
+
+API_BASE_URL = "https://api.smith.langchain.com/api/v1"
+
+class LangSmithAPI:
+
+    def __init__(self):
+        self.base_url = API_BASE_URL
+        self.headers = {
+            'X-API-Key': os.environ.get("LANGCHAIN_API_KEY")
+        }
+
+        self.getSessionId()
+
+    def getSessionId(self):
+
+        payload = {}
+
+        response = requests.request("GET", f"{self.base_url}/sessions", headers=self.headers, data=payload)
+        response_json = response.json()
+
+        for session in response_json:
+            session_id = session.get('id', None)
+            name = session.get('name', 'Unnamed')
+            if( name == os.environ.get("LANGCHAIN_PROJECT") ):
+                self.session_id = session_id
+
+
+    def getLastTraces(self, id_request: str = ''):
+        payload = {
+            "session": [self.session_id],
+            "run_type":"chain",
+            "select": [
+                "id","trace_id","name","run_type","extra","start_time","inputs",
+                "total_tokens","prompt_tokens","completion_tokens",
+                "total_cost","prompt_cost","completion_cost","first_token_time"
+            ],
+            "order": "desc",
+            "limit": 100
+        }
+
+        # Make request
+        response = requests.request("POST", f"{self.base_url}/runs/query", headers=self.headers, data=json.dumps(payload))
+        response_json = response.json()
+           
+        # filters for the traces that have the same id_request
+        # filtered_traces = [trace for trace in response_json['runs'] if 'extra' in trace and 'metadata' in trace['extra'] and 'id_request' in trace['extra']['metadata'] and trace['extra']['metadata']['id_request'] == id_request]
+        filtered_traces = [trace for trace in response_json['runs'] if trace.get('parent_run_ids') is not None and trace['parent_run_ids'] == []]
+        
+        return filtered_traces
+        
+    
+    def getDailyMessages(self, id_request: str = ""):
+
+        trace = self.getLastTraces(id_request)
+        
+        daily_counter = 0
+        
+        for t in trace:
+            start_time = datetime.fromisoformat(t['start_time'])
+            if start_time.date() == datetime.today().date():
+                daily_counter += 1
+
+        return daily_counter

app.py

@@ -5,6 +5,7 @@ from dotenv import load_dotenv
 from rag import Rag
 from vectore_store.PineconeConnector import PineconeConnector
 from vectore_store.VectoreStoreManager import VectoreStoreManager
+from subscription.index import get_usages, is_usage_limit_reached
 
 from util import getYamlConfig
 
@@ -49,6 +50,16 @@ def main():
 
     st.logo(LOGO)
     st.title(GROUP_NAME)
+    
+    usages = get_usages()
+    reached = is_usage_limit_reached()
+
+    st.sidebar.html(f"<p style='margin:0;text-align:center;font-weight:bold;font-size:24px;'>{usages['count']} / {usages['daily_limit']}</p>")
+    if reached:
+        st.sidebar.html(f"<p style='text-align:center;font-style:italic;'>Vous avez atteint la limite d'utilisation<p>")
+    else:
+        st.sidebar.html(f"<p style='text-align:center;'>Nombre de messages envoyés aujourd'hui<p>")
+    
 
     saved_documents = st.Page("pages/persistent_documents.py", title="Communs", icon="🗃️")
     documents = st.Page("pages/documents.py", title="Vos documents", icon="📂")

pages/chatbot.py

@@ -4,6 +4,7 @@ from langchain.prompts import PromptTemplate
 from model import selector
 from util import getYamlConfig
 from st_copy_to_clipboard import st_copy_to_clipboard
+from subscription.index import is_usage_limit_reached
 
 def display_messages():
 
@@ -123,8 +124,13 @@ def page():
     # Displaying messages
     display_messages()
 
+    reached = is_usage_limit_reached()
+    message = ""
+    if reached:
+        message = "Vous avez atteint la limite quotidienne de messages."
+
+    user_query = st.chat_input(message, disabled=reached)
 
-    user_query = st.chat_input("")
     if user_query is not None and user_query != "":
 
         st.session_state["chat_history"].append(HumanMessage(content=user_query))

subscription/index.py

@@ -0,0 +1,20 @@
+import os
+from apis.LangSmith import LangSmithAPI
+from dotenv import load_dotenv
+
+load_dotenv()
+
+langsmith = LangSmithAPI()
+count = langsmith.getDailyMessages()
+
+# Convert DAILY_LIMIT to int with a default value of 5
+DAILY_LIMIT = int(os.getenv("DAILY_MESSAGES_LIMIT", "5"))
+
+def is_usage_limit_reached():
+    return count >= DAILY_LIMIT
+
+def get_usages():
+    return {
+        "count": count,
+        "daily_limit": DAILY_LIMIT
+    }