give every user 100,000 free tokens

Author: CuriousDima

.env.template

@@ -1,5 +1,5 @@
-TELEGRAM_BOT_TOKEN=<your-telegram-bot-token-value-here>
 DB_URI=<choose-a-path-for-sqlite-db-file>
+TELEGRAM_BOT_TOKEN=<your-telegram-bot-token-value-here>
 OPENAI_API_KEY=<your-openai-key-here>
 # https://console.groq.com/keys
 GROQ_API_KEY=<your-groq-key-here>

telegram-bot/bot.py

@@ -14,6 +14,7 @@
 )
 
 from llm_client import LLMClient, Model, Provider
+from db_client import DBClient, Account
 
 
 logging.basicConfig(
@@ -25,6 +26,7 @@
 
 # Environment variables names
 _TELEGRAM_BOT_TOKEN_VAR_NAME = "TELEGRAM_BOT_TOKEN"
+_DB_URI_VAR_NAME = "DB_URI"
 
 _START_MESSAGE = (
     "Hi there! I'm a bot designed to assist you in rephrasing your text into polished English. "
@@ -38,30 +40,51 @@ async def start_command(update: Update, _: ContextTypes.DEFAULT_TYPE) -> None:
 
 
 async def rewrite(
-    update: Update, context: ContextTypes.DEFAULT_TYPE, llm_client: LLMClient
+    update: Update,
+    context: ContextTypes.DEFAULT_TYPE,
+    db_client: DBClient,
+    llm_client: LLMClient,
 ) -> None:
+    # Handle the message
     input_message = None
     if update.message is not None:
         input_message = update.message.text
     elif update.edited_message is not None:
         input_message = update.edited_message.text
     assert input_message is not None, "No message to rewrite."
-    rewritten_text = llm_client.rewrite(input_message)
+    # Handle the user
+    user_id = update.message.from_user.id
+    username = update.message.from_user.username
+    account = db_client.get_or_create_account(user_id=user_id, username=username)
+    # Check if the user has run out of tokens.
+    # If the user is a friend, they have an unlimited token balance. ;)
+    if not account.is_friend and account.tokens_balance <= 0:
+        await context.bot.send_message(
+            chat_id=update.effective_chat.id,
+            text="You have run out of tokens. 🥲\n Please contact the bot owner to get more.",
+        )
+        return
+    # Rewrite the message. Do not touch the user's token balance if they are a friend.
+    rewritten_text, num_tokens = llm_client.rewrite(input_message)
+    if not account.is_friend:
+        db_client.decrease_token_balance(account=account, num_tokens=num_tokens)
     await context.bot.send_message(
         chat_id=update.effective_chat.id,
         text=rewritten_text,
     )
 
 
 if __name__ == "__main__":
+    db_client = DBClient(db_url=os.getenv(_DB_URI_VAR_NAME))
+    llm_client = LLMClient(provider=Provider.GROQ, model=Model.GEMMA)
+
     app = ApplicationBuilder().token(os.getenv(_TELEGRAM_BOT_TOKEN_VAR_NAME)).build()
 
     app.add_handler(CommandHandler("start", start_command))
-
-    llm_client = LLMClient(provider=Provider.GROQ, model=Model.GEMMA)
     app.add_handler(
         MessageHandler(
-            filters.TEXT & (~filters.COMMAND), partial(rewrite, llm_client=llm_client)
+            filters.TEXT & (~filters.COMMAND),
+            partial(rewrite, db_client=db_client, llm_client=llm_client),
         )
     )
 

telegram-bot/db_client.py

@@ -0,0 +1,45 @@
+from sqlalchemy import BigInteger
+from sqlmodel import create_engine, SQLModel, Session, select, Field
+
+from llm_client import Model, Provider
+
+_NUM_TOKENS_DEFAULT = 100_000
+
+
+class Account(SQLModel, table=True):
+    id: int | None = Field(default=None, primary_key=True)
+    user_id: str = Field(index=True, sa_type=BigInteger)  # Telegram user ID
+    username: str | None
+    provider: Provider = Field(default=Provider.GROQ)
+    model: Model = Field(default=Model.GEMMA)
+    tokens_balance: int = Field(default=_NUM_TOKENS_DEFAULT)
+    # Whether the user is a friend of the bot owner.
+    # This is used to give the user an unlimited token balance.
+    is_friend: bool = Field(default=False)
+
+
+class DBClient:
+    def __init__(self, db_url: str, echo=True) -> None:
+        self.engine = create_engine(db_url, echo=echo)
+        SQLModel.metadata.create_all(self.engine)
+
+    def __del__(self) -> None:
+        self.engine.dispose()
+
+    def get_or_create_account(self, user_id: str, username=None) -> Account:
+        with Session(self.engine) as session:
+            statement = select(Account).filter(Account.user_id == user_id)
+            account = session.exec(statement).one_or_none()
+            if account is None:
+                account = Account(user_id=user_id, username=username)
+                session.add(account)
+                session.commit()
+                session.refresh(account)
+            return account
+
+    def decrease_token_balance(self, account: Account, num_tokens: int) -> None:
+        with Session(self.engine) as session:
+            account.tokens_balance -= num_tokens
+            session.add(account)
+            session.commit()
+            session.refresh(account)

telegram-bot/llm_client.py

@@ -1,5 +1,5 @@
-from enum import Enum, auto
-from typing import Callable, Union
+from enum import Enum
+from typing import Callable, Union, Tuple
 import os
 
 from langchain_core.messages import SystemMessage
@@ -33,15 +33,15 @@
 
 
 class Provider(Enum):
-    OPENAI = auto()
-    GROQ = auto()
+    GROQ = "groq"
+    OPENAI = "openai"
 
 
 class Model(Enum):
+    GEMMA = "gemma-7b-it"
     GPT3 = "gpt-3.5-turbo"
     GPT4 = "gpt-4.0-turbo"
     MIXTRAL = "mixtral-8x7b-32768"
-    GEMMA = "gemma-7b-it"
 
 
 # OpenAI can be used as a provider with the following models: GPT3, GPT4.
@@ -82,14 +82,18 @@ def get_chat(
 
 class LLMClient:
     def __init__(
-        self, provider: Provider, model: Model, temperature: float = 0.5
+        self, provider: Provider, model: Model, temperature: float = 0.35
     ) -> None:
         self.chat = get_chat(provider, model)(model, temperature)
 
-    def rewrite(self, text: str) -> str:
+    def rewrite(self, text: str) -> Tuple[str, int]:
         output_parser = StrOutputParser()
-        chain = _REWRITE_PROMPT | self.chat | output_parser
-        return chain.invoke(
+        chain = _REWRITE_PROMPT | self.chat
+        output = chain.invoke(
             {"text": text},
             config={"callbacks": [ConsoleCallbackHandler()]},
         )
+        return (
+            output_parser.invoke(output),
+            output.response_metadata["token_usage"]["total_tokens"],
+        )