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.
 

101 lines
3.2 KiB

import os
from langchain.agents import Tool, initialize_agent, AgentType
from dotenv import load_dotenv
from tools import DownloadListTool, QBitSearchTool, DownloadTorrentTool
from langchain_community.tools import DuckDuckGoSearchRun
from langchain.memory import ConversationBufferMemory
from langchain.chat_models import init_chat_model
import gradio as gr
# Load environment variables
load_dotenv()
def create_agent():
# Initialize the language model
llm = init_chat_model("gpt-4o-mini", model_provider="openai")
# Initialize memory
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
# Initialize search tool
search_tool = DuckDuckGoSearchRun()
# Function to force DuckDuckGo for specific search types
def forced_duckduckgo_search(query: str) -> str:
"""Use DuckDuckGo to search for specific information."""
return search_tool.run(query)
# Initialize tools
tools = [
DownloadListTool(),
QBitSearchTool(),
DownloadTorrentTool(),
Tool(
name="ForcedDuckDuckGoSearch",
func=forced_duckduckgo_search,
description="Use this tool when you need to find specific information about movies, TV shows. Input should be a search query including the keyword 'imdb'.",
)
]
# 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="qBittorrent AI Agent") as interface:
gr.Markdown("# qBittorrent AI Agent")
gr.Markdown("Ask questions about downloads, search for content, 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"],
title="qBittorrent Assistant"
)
# Launch the interface
interface.launch(share=False)
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()