qbit-agent/main.py

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()