ia-subvention / apis /LangSmith.py
Reset daily message count at the start of a new day
18107fd
raw
history blame
2.9 kB
import requests
import json
import os
from dotenv import load_dotenv
from datetime import datetime, timedelta
load_dotenv()
API_BASE_URL = "https://api.smith.langchain.com/api/v1"
class LangSmithAPI:
_daily_count = 0
_last_count_date = None
def __init__(self):
self.base_url = API_BASE_URL
self.headers = {
'X-API-Key': os.environ.get("LANGCHAIN_API_KEY")
}
self.getSessionId()
if self._last_count_date != datetime.today().date():
self._daily_count = 0
self._last_count_date = datetime.today().date()
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)
current_date = datetime.today().date()
# current_date = current_date - timedelta(days=1) # Yesterday
# Reset counter if it's a new day
if self._last_count_date != current_date:
self._daily_count = 0
self._last_count_date = current_date
# Count only new messages from the current day
for t in trace:
start_time = datetime.fromisoformat(t['start_time'])
if start_time.date() == current_date:
self._daily_count += 1
return self._daily_count