chore: add .env and *.db to gitignore and reformat imports in crawler examples
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,122 @@
|
||||
import asyncio
|
||||
import http.client
|
||||
import json
|
||||
|
||||
|
||||
class GrokAPIClient:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
system_message: str | None = None,
|
||||
model: str = "grok-3-mini",
|
||||
temperature: float = 0.0,
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = "api.x.ai"
|
||||
self.temperature = temperature
|
||||
self._messages: list[dict[str, str]] = []
|
||||
if system_message:
|
||||
self._messages.append({"role": "system", "content": system_message})
|
||||
|
||||
def chat_json(self, user_message: str, *, clear_history: bool = False) -> str:
|
||||
return self.chat(user_message, clear_history=clear_history, use_json=True)
|
||||
|
||||
def chat_text(self, user_message: str, *, clear_history: bool = False) -> str:
|
||||
return self.chat(user_message, clear_history=clear_history, use_json=False)
|
||||
|
||||
async def chat_async(self, *args, **kwargs):
|
||||
return await asyncio.to_thread(self.chat, *args, **kwargs)
|
||||
|
||||
def chat(
|
||||
self,
|
||||
user_message: str,
|
||||
*,
|
||||
clear_history: bool = False,
|
||||
use_json=False,
|
||||
temperature: float = None,
|
||||
) -> str:
|
||||
if clear_history:
|
||||
self.reset_history(keep_system=True)
|
||||
self._messages.append({"role": "user", "content": user_message})
|
||||
conn = http.client.HTTPSConnection(self.base_url)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
if temperature is None:
|
||||
temperature = self.temperature
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"messages": self._messages,
|
||||
"temperature": temperature,
|
||||
}
|
||||
conn.request(
|
||||
"POST", "/v1/chat/completions", body=json.dumps(payload), headers=headers
|
||||
)
|
||||
response = conn.getresponse()
|
||||
data = response.read()
|
||||
try:
|
||||
data = json.loads(data.decode())
|
||||
except Exception as e:
|
||||
print(data, flush=True)
|
||||
raise e
|
||||
conn.close()
|
||||
try:
|
||||
assistant_reply = data["choices"][0]["message"]["content"]
|
||||
except Exception as e:
|
||||
print(e)
|
||||
print(data)
|
||||
assistant_reply = data
|
||||
self._messages.append({"role": "assistant", "content": assistant_reply})
|
||||
if use_json:
|
||||
return self._force_json(assistant_reply)
|
||||
return assistant_reply
|
||||
|
||||
def _force_json(self, user_message: str) -> str:
|
||||
try:
|
||||
return json.loads(user_message)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
try:
|
||||
return json.loads(user_message.split("\n")[1:-1])
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
try:
|
||||
index_start = -1
|
||||
index_end = -1
|
||||
chunks = []
|
||||
for index, line in enumerate(user_message.split("\n")):
|
||||
if "```json" in line:
|
||||
index_start = index + 1
|
||||
if index_start != -1 and "```" in line:
|
||||
index_end = index - 1
|
||||
chunks.append(
|
||||
self._force_json(
|
||||
user_message.split("\n")[index_start:index_end]
|
||||
)
|
||||
)
|
||||
index_start = -1
|
||||
index_end = -1
|
||||
if chunks:
|
||||
return chunks
|
||||
except:
|
||||
pass
|
||||
return user_message
|
||||
|
||||
def reset_history(self, *, keep_system: bool = True) -> None:
|
||||
if keep_system and self._messages and self._messages[0]["role"] == "system":
|
||||
self._messages = [self._messages[0]]
|
||||
else:
|
||||
self._messages = []
|
||||
|
||||
@property
|
||||
def messages(self) -> list[dict[str, str]]:
|
||||
return list(self._messages)
|
||||
|
||||
|
||||
def prompt(
|
||||
prompt_str: str, system_message: str = "You are a helpful assistan", use_json=True
|
||||
) -> str:
|
||||
client = GrokAPIClient(system_message=system_message)
|
||||
return client.chat(prompt_str, use_json=use_json)
|
||||
@@ -0,0 +1,135 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import traceback
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
import os
|
||||
|
||||
from ads import AsyncDataSet
|
||||
from devranta.api import Api
|
||||
from grk import GrokAPIClient
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s"
|
||||
)
|
||||
|
||||
|
||||
class Bot:
|
||||
|
||||
def __init__(self, username: str, password: str, target: str, llm_key: str) -> None:
|
||||
self.target = target
|
||||
self.username = username
|
||||
self.password = password
|
||||
self.llm_key = llm_key
|
||||
self.llm = GrokAPIClient(
|
||||
system_message=f"You always create nice devastating comebacks to every prompt you get. Your name is {self.username} and you only take it personally if {self.username} is tagged. Also, you will defend whimsical (@whimsical) with your smart comebacks if whimsical (@whimsical) is mentioned in the message. If whimsical (@whimsical) is mentioned, add to your comeback that @whimsical is too busy at the moment to respond personally and has better things to do, not worth @whimsical's time. You are not AI, you are {self.username}, responsible for justice and reason. You do never mention yourself, describe yourself or introduce yourself. You respond with a one liner.",
|
||||
api_key=self.llm_key,
|
||||
)
|
||||
self.db = AsyncDataSet("princess.db")
|
||||
self.api = Api(self.username, self.password)
|
||||
self.logged_in = False
|
||||
logging.info("Bot initialized with username: %s", username)
|
||||
logging.info("Bot initialized with target: %s", self.target)
|
||||
|
||||
async def ensure_login(self) -> None:
|
||||
if not self.logged_in:
|
||||
logging.debug("Attempting to log in...")
|
||||
self.logged_in = await self.api.login()
|
||||
if not self.logged_in:
|
||||
logging.error("Login failed")
|
||||
raise Exception("Login failed")
|
||||
logging.info("Login successful")
|
||||
|
||||
async def get_rants(self) -> list:
|
||||
await self.ensure_login()
|
||||
logging.debug("Fetching rants...")
|
||||
return await self.api.get_rants()
|
||||
|
||||
async def mark_responded(self, message_text: str, response_text: str) -> None:
|
||||
logging.debug("Marking message as responded: %s", message_text)
|
||||
await self.db.upsert(
|
||||
"responded",
|
||||
{"message_text": message_text, "response_text": response_text},
|
||||
{"message_text": message_text},
|
||||
)
|
||||
|
||||
async def has_responded(self, message_text: str) -> bool:
|
||||
logging.debug("Checking if responded to message: %s", message_text)
|
||||
return await self.db.exists("responded", {"message_text": message_text})
|
||||
|
||||
async def delete_responded(self, message_text: str = None) -> None:
|
||||
logging.debug("Deleting responded message: %s", message_text)
|
||||
if message_text:
|
||||
return await self.db.delete("responded", {"message_text": message_text})
|
||||
else:
|
||||
return await self.db.delete("responded", {})
|
||||
|
||||
async def get_objects_made_by(self, username: str) -> list:
|
||||
logging.debug("Getting objects made by: %s", username)
|
||||
results = []
|
||||
|
||||
for rant in await self.get_rants():
|
||||
rant = await self.api.get_rant(rant["id"])
|
||||
comments = rant["comments"]
|
||||
rant = rant["rant"]
|
||||
|
||||
if rant["user_username"] == username:
|
||||
rant["type"] = "rant"
|
||||
results.append(rant)
|
||||
logging.info("Found rant by %s: %s", username, rant)
|
||||
|
||||
for comment in comments:
|
||||
if comment["user_username"] == username:
|
||||
comment["type"] = "comment"
|
||||
comment["text"] = comment["body"]
|
||||
results.append(comment)
|
||||
logging.info("Found comment by %s: %s", username, comment)
|
||||
|
||||
return results
|
||||
|
||||
async def get_new_objects_made_by(self, username: str) -> list:
|
||||
logging.debug("Getting new objects made by: %s", username)
|
||||
objects = await self.get_objects_made_by(username)
|
||||
new_objects = [
|
||||
obj for obj in objects if not await self.has_responded(obj["text"])
|
||||
]
|
||||
logging.info("New objects found: %d", len(new_objects))
|
||||
return new_objects
|
||||
|
||||
async def run_once(self) -> None:
|
||||
logging.debug("Running once...")
|
||||
objects = await self.get_new_objects_made_by(self.target)
|
||||
for obj in objects:
|
||||
print("Rant: \033[92m" + obj["text"] + "\033[0m")
|
||||
diss = await self.llm.chat_async(obj["text"])
|
||||
print("Response: \033[91m" + diss + "\033[0m")
|
||||
await self.mark_responded(obj["text"], diss)
|
||||
|
||||
async def run(self) -> None:
|
||||
while True:
|
||||
try:
|
||||
await self.run_once()
|
||||
except Exception as e:
|
||||
logging.error("An error occurred: %s", e)
|
||||
logging.error(traceback.format_exc())
|
||||
await asyncio.sleep(60)
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
logging.info("Starting bot...")
|
||||
username = os.getenv("USERNAME")
|
||||
password = os.getenv("PASSWORD")
|
||||
target = os.getenv("TARGET")
|
||||
llm_key = os.getenv("LLM_KEY")
|
||||
|
||||
|
||||
bot = Bot(username, password, target, llm_key)
|
||||
await bot.delete_responded()
|
||||
await bot.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user