You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 

107 lines
4.3 KiB

import os
from langchain.agents import Tool, initialize_agent, AgentType
from dotenv import load_dotenv
from tools.qbit import QbitDownloadListTool, QbitSearchTool, QbitDownloadTorrentTool
from tools.search import MediaInfoSearchTool, MoviesAdviceSearchTool
from langchain.memory import ConversationBufferMemory
from langchain.chat_models import init_chat_model
import gradio as gr
import datetime
# Load environment variables
load_dotenv()
def create_agent():
# Initialize the language model
#llm = init_chat_model("gpt-4.1-mini", model_provider="openai")
llm = init_chat_model("qwen2.5-coder:14b", model_provider="ollama", temperature=0)
# Create system message with current time and other info
current_time = datetime.datetime.now()
# Initialize memory
memory = ConversationBufferMemory(
memory_key="chat_history",
return_messages=True,
human_prefix="User",
ai_prefix="Assistant"
)
memory.chat_memory.add_message(f"SYSTEM: today is {current_time.strftime('%Y-%m-%d')}")
memory.chat_memory.add_message(f"SYSTEM:from now on when User ask for movie or tv series suggestion reply with a numbered markdown list with a brief description of each title")
memory.chat_memory.add_message(f"SYSTEM:from now on when list torrents show seeds number and MAGNET LINK (USING A MARKDOWN LINK WITH TEXT 'Magnet link 🧲')")
memory.chat_memory.add_message(f"SYSTEM:from now on, when show downloads list show a clean and nice markdown format with name and the most important information, \
also add near it an emoji of progress of the download to represent how complete it is and set a coloured bullet emoji after status of torrent status, for example blue for uploading, green for downloading, \
red for error, yellow for paused, and grey for completed")
memory.chat_memory.add_message(f"SYSTEM: from now on, when user ask for downolad NEVER start a qbittorrent download if user hasn't viewed the list of torrents first, \
and choosed one of them")
# Initialize tools
tools = [
QbitDownloadListTool(),
QbitSearchTool(),
QbitDownloadTorrentTool(),
MoviesAdviceSearchTool(),
MediaInfoSearchTool()
]
# Initialize the agent with memory
agent = initialize_agent(
tools,
llm,
agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
verbose=True,
memory=memory
)
return agent
def process_query(message, history):
try:
# Create agent if it doesn't exist
if not hasattr(process_query, "agent"):
process_query.agent = create_agent()
# Run the agent with the user's message
response = process_query.agent.run(message)
return response
except Exception as e:
return f"Error: {str(e)}"
def main():
print("Starting qBittorrent AI Agent...")
# Create Gradio interface
with gr.Blocks(title="qbit-agent") as interface:
gr.Markdown("# qbit-agent")
gr.Markdown("### Made by Matteo with hate and piracy 💀")
gr.Markdown("Ask about downloads, search for content (and torrent), or get recommendations.")
chatbot = gr.ChatInterface(
process_query,
examples=["Find me the latest sci-fi movies",
"What are the top TV shows from 2023?",
"Download Interstellar in 1080p",
"Show me my current downloads",
"What is The Matrix",
"Get me a list of horror movies"],
)
# Launch the interface
interface.launch(share=True)
def cli_main():
print("Starting qBittorrent AI Agent in CLI mode...")
agent = create_agent()
while True:
user_input = input("\nEnter your question (or 'quit' to exit): ")
if user_input.lower() in ['quit', 'exit']:
break
try:
response = agent.run(user_input)
print(response)
except Exception as e:
print(f"Error: {str(e)}")
if __name__ == "__main__":
# Use main() for Gradio interface or cli_main() for command-line interface
main()
# Uncomment the line below to use CLI instead
# cli_main()