PikoAi 0.1.16__py3-none-any.whl → 0.1.18__py3-none-any.whl

Agents/Executor/executor.py

@@ -9,16 +9,7 @@ from Utils.ter_interface import TerminalInterface
 from Utils.executor_utils import parse_tool_call
 from Agents.Executor.prompts import get_executor_prompt # Import prompts
 
-from typing import Optional
-from mistralai.models.sdkerror import SDKError # This might be an issue if LiteLLM doesn't use SDKError
-                                              # LiteLLM maps exceptions to OpenAI exceptions.
-                                              # We'll keep it for now and see if errors arise during testing.
-# from Env import python_executor # Will be replaced by BaseEnv
-# from Env.shell import ShellExecutor # Will be replaced by BaseEnv
-from Env.base_env import create_environment, BaseEnv # Added
-from Env import python_executor # Keep for type hint in the old execute method if needed, or remove if execute is fully removed
-from llm_interface.llm import LiteLLMInterface # Import LiteLLMInterface
-
+from llm_interface.llm import LiteLLMInterface # Import LiteLLMInterfacea
 from Tools import tool_manager
 
 class RateLimiter:
@@ -38,7 +29,7 @@ class executor:
     def __init__(self, user_prompt, max_iter=10):
         self.user_prompt = user_prompt
         self.max_iter = max_iter
-        self.rate_limiter = RateLimiter(wait_time=5.0, max_retries=3)
+        self.rate_limiter = RateLimiter(wait_time=3.0, max_retries=3)
         self.executor_prompt_init()  # Update system_prompt
         # self.python_executor = python_executor.PythonExecutor()  # Initialize PythonExecutor
         # self.shell_executor = ShellExecutor() # Initialize ShellExecutor
@@ -89,16 +80,10 @@ class executor:
 
             except Exception as e: # Catching generic Exception as LiteLLM maps to OpenAI exceptions
                 # Check if the error message contains "429" for rate limiting
-                if "429" in str(e) and retries < self.rate_limiter.max_retries:
+                if retries < self.rate_limiter.max_retries:
                     retries += 1
                     print(f"\nRate limit error detected. Waiting {self.rate_limiter.wait_time} seconds before retry {retries}/{self.rate_limiter.max_retries}")
                     time.sleep(self.rate_limiter.wait_time)
-                # Check if the error is an SDKError (though less likely with LiteLLM directly)
-                # or if it's any other exception that we should retry or raise.
-                elif isinstance(e, SDKError) and "429" in str(e) and retries < self.rate_limiter.max_retries: # Added SDKError check just in case
-                    retries += 1
-                    print(f"\nRate limit exceeded (SDKError). Waiting {self.rate_limiter.wait_time} seconds before retry {retries}/{self.rate_limiter.max_retries}")
-                    time.sleep(self.rate_limiter.wait_time)
                 else:
                     print(f"\nError occurred during inference: {str(e)}")
                     # You might want to log the full traceback here for debugging

Agents/Executor/prompts.py

@@ -33,9 +33,7 @@ You must break down the user's goal into smaller steps and perform one action at
     }}
   }}
   <<END_TOOL_CALL>>
-- **Code Execution**: Write Python code when no tool is suitable or when custom logic is needed.
-the code written will be executed immediately and not saved. 
-- **Direct Response**: Provide a direct answer if the task doesn't require tools or code.
+- **Direct Response**: Provide a direct answer if the task doesn't require tool calling
 
 
 These are the things that you learned from the mistakes you made earlier :
@@ -46,13 +44,13 @@ These are the things that you learned from the mistakes you made earlier :
   - Don't execute dangerous commands like rm -rf * or access sensitive files
   - If you are stuck, have tried to fix an issue (e.g., a linter error) multiple times (e.g., 3 times) without success, or need clarification, ask the USER for input. Explain the situation clearly.
   - Upon creating anything (like a new project, website, data analysis png) always show the output.You can do this by executing shell commands.
-  - the python/shell code execution in tool call will be executed immediately and output will be shown. it wont be saved.
+  - the python/shell code execution through tool call will be executed immediately and output will be shown. it wont be saved.
 
 
 ** Important **
-- Perform only one action per step (either a single tool call or a single code execution).
-- Always evaluate the output of each action before deciding the next step.
+- You can only perform one tool call at a time.
+- Always evaluate the output of the tool call before deciding the next step.
 - Continue performing actions until the user's goal is fully achieved. Only then, include 'TASK_DONE' in your response if that is the required signal for completion.
-- Do not end the task immediately after a tool call or code execution without evaluating its output.
+- Do not end the task immediately after a tool call without evaluating its output.
 
 """

Tools/tool_manager.py

@@ -77,7 +77,7 @@ def call_tool(tool_name, tool_input):
         # Pass the tool_input dictionary as kwargs to the tool function
         return tools_function_map[tool_name](**tool_input)
     else:
-        raise ValueError(f"Tool '{tool_name}' not found. Check the tools available in the tool directory")
+        return f"This tool is invalid. Please check the tools available in the tool directory"
 
 tools_function_map = {
     "web_loader": load_data,

Tools/web_search.py

@@ -1,8 +1,14 @@
+import os
+from dotenv import load_dotenv
 from duckduckgo_search import DDGS
+import requests
+
+load_dotenv()
 
 def web_search(max_results: int = 10, **kwargs) -> str:
     """
     Performs a DuckDuckGo web search based on your query (think a Google search) then returns the top search results.
+    Falls back to SerpAPI (Google) if DuckDuckGo fails or returns no results.
 
     Args:
         query (str): The search query to perform.
@@ -13,18 +19,40 @@ def web_search(max_results: int = 10, **kwargs) -> str:
         str: Formatted string containing search results.
 
     Raises:
-        ImportError: If the duckduckgo_search package is not installed.
+        ImportError: If neither duckduckgo_search nor SerpAPI is available.
         Exception: If no results are found for the given query.
     """
+    query = kwargs['query']
+    # Try DuckDuckGo first
     try:
         ddgs = DDGS()
-    except ImportError as e:
-        raise ImportError("You must install package `duckduckgo_search` to run this function: for instance run `pip install duckduckgo-search`."
-        ) from e
-    query = kwargs['query']
-    results = ddgs.text(query, max_results=max_results)
-    if len(results) == 0:
-        raise Exception("No results found! Try a less restrictive/shorter query.")
+        results = ddgs.text(query, max_results=max_results)
+        if results and len(results) > 0:
+            postprocessed_results = [f"[{result['title']}]({result['href']})\n{result['body']}" for result in results]          
+            return "## Search Results (DuckDuckGo)\n\n" + "\n\n".join(postprocessed_results)
+    except Exception:
+        pass  # Will try SerpAPI fallback
 
-    postprocessed_results = [f"[{result['title']}]({result['href']})\n{result['body']}" for result in results]
-    return "## Search Results\n\n" + "\n\n".join(postprocessed_results)
+    # Fallback to SerpAPI (Google)
+    SERP_API_KEY = os.getenv("SERP_API_KEY")
+    if not SERP_API_KEY:
+        raise ImportError("DuckDuckGo search failed and SERP_API_KEY is not set for SerpAPI fallback.")
+    url = "https://serpapi.com/search"
+    params = {
+        "q": query,
+        "api_key": SERP_API_KEY,
+        "engine": "google"
+    }
+    try:
+        response = requests.get(url, params=params)
+        data = response.json()
+        results = data.get("organic_results", [])
+        if not results:
+            raise Exception("No results found from SerpAPI either!")
+        simplified_results = [
+            f"[{result.get('title')}]({result.get('link')})\n{result.get('snippet', '')}"
+            for result in results[:max_results]
+        ]
+        return "## Search Results (Google via SerpAPI)\n\n" + "\n\n".join(simplified_results)
+    except Exception as e:
+        raise Exception(f"No results found from DuckDuckGo or SerpAPI. Last error: {e}")

Utils/ter_interface.py

@@ -68,7 +68,8 @@ class TerminalInterface:
 
             # Handle tool call closing delimiter - be more flexible with whitespace
             elif "<<END_TOOL_CALL>>" in line_stripped:
-                self.console.print(Syntax('{"status": "end_tool_call"}', "json", theme="monokai", line_numbers=False))
+                self.console.print(self.tool_call_buffer)
+                # self.console.print(Syntax('{"status": "end_tool_call"}', "json", theme="monokai", line_numbers=False))
                 self.console.print("[bold cyan]--------------------------------[/bold cyan]")
                 self.inside_tool_call = False
                 self.tool_call_buffer = ""

cli.py

@@ -32,6 +32,10 @@ AVAILABLE_MODELS = {
         "anthropic/claude-3-opus-20240229",
         "anthropic/claude-3-sonnet-20240229",
         "anthropic/claude-3-haiku-20240307"
+    ],
+    "gemini": [
+        "gemini/gemini-2.0-flash",
+        "gemini/gemini-2.5-flash-preview-05-20"
     ]
 }
 
@@ -40,323 +44,257 @@ API_KEYS = {
     "openai": "OPENAI_API_KEY",
     "mistral": "MISTRAL_API_KEY",
     "groq": "GROQ_API_KEY",
-    "anthropic": "ANTHROPIC_API_KEY"
+    "anthropic": "ANTHROPIC_API_KEY",
+    "gemini": "GEMINI_API_KEY"
 }
 
+# --- Utility Functions ---
+
+def clear_terminal():
+    """Clear the terminal screen."""
+    os.system('cls' if os.name == 'nt' else 'clear')
+
 def get_provider_from_model_name(model_name: str) -> str:
     """Extracts the provider from a litellm model string (e.g., 'openai/gpt-4o' -> 'openai')."""
     if not model_name or '/' not in model_name:
-        # Fallback or error handling if model_name is not in expected format
-        # For now, try to return the model_name itself if it doesn't contain '/',
-        # as it might be a provider name already or an old format.
-        # This case should ideally be handled based on how robust the system needs to be.
         print(f"Warning: Model name '{model_name}' may not be in 'provider/model' format. Attempting to use as provider.")
         return model_name 
     return model_name.split('/')[0]
 
-def clear_terminal():
-    """Clear the terminal screen"""
-    os.system('cls' if os.name == 'nt' else 'clear')
-
-def ensure_api_key(provider):
-    """Ensure API key exists for the given provider"""
-    env_path = os.path.join(os.path.dirname(__file__), '../.env')
-    env_var = API_KEYS.get(provider)
-    if not env_var:
-        raise ValueError(f"Unknown provider: {provider}")
-    
-    # Force reload of .env file
-    if os.path.exists(env_path):
-        load_dotenv(env_path, override=True)
-    
-    # Check if API key exists in environment
-    api_key = os.getenv(env_var)
-    
-    if not api_key:
-        # Ask for API key
-        questions = [
-            inquirer.Text('api_key',
-                message=f"Enter your {provider.upper()} API key",
-                validate=lambda _, x: len(x.strip()) > 0
-            )
-        ]
-        api_key = inquirer.prompt(questions)['api_key']
-        clear_terminal()
-        
-        # Save to .env file
-        if not os.path.exists(env_path):
-            with open(env_path, 'w') as f:
-                f.write(f"{env_var}={api_key}\n")
-        else:
-            # Read existing .env file
-            with open(env_path, 'r') as f:
-                lines = f.readlines()
-            
-            # Check if key already exists
-            key_exists = False
-            for i, line in enumerate(lines):
-                if line.strip().startswith(f"{env_var}=") or line.strip().startswith(f"#{env_var}="):
-                    lines[i] = f"{env_var}={api_key}\n"
-                    key_exists = True
-                    break
-            
-            # If key doesn't exist, append it
-            if not key_exists:
-                lines.append(f"{env_var}={api_key}\n")
-            
-            # Write back to .env file
-            with open(env_path, 'w') as f:
-                f.writelines(lines)
-        
-        # Reload environment with new key
-        load_dotenv(env_path, override=True)
-    
-    return api_key
+# --- Configuration Management ---
 
-def ensure_config_exists():
-    """Ensure config.json exists and has required fields"""
-    config_path = os.path.join(os.path.dirname(__file__), '../config.json')
-    config = None
-    
+def load_config(config_path: str) -> dict:
+    """Load (or create) config.json and return its contents as a dict. If config.json does not exist, copy from config.example.json (or create a default) and update working_directory to os.getcwd()."""
     if not os.path.exists(config_path):
-        # Copy config from example if it doesn't exist
         example_path = os.path.join(os.path.dirname(__file__), '../config.example.json')
         if os.path.exists(example_path):
             shutil.copy2(example_path, config_path)
             with open(config_path, 'r') as f:
                 config = json.load(f)
         else:
-            # Create a default config if example is missing
-            config = {
-                "working_directory": os.getcwd(),
-                "llm_provider": None, # Will store the provider part, e.g., "openai"
-                "model_name": None    # Will store the full litellm string, e.g., "openai/gpt-4o"
-            }
+            config = { "working_directory": os.getcwd(), "llm_provider": None, "model_name": None }
     else:
-        # Read existing config
         with open(config_path, 'r') as f:
             try:
                 config = json.load(f)
             except json.JSONDecodeError:
                 print("Error reading config.json. File might be corrupted. Re-creating default.")
-                config = {
-                    "working_directory": os.getcwd(),
-                    "llm_provider": None,
-                    "model_name": None
-                }
+                config = { "working_directory": os.getcwd(), "llm_provider": None, "model_name": None }
+    # Always update working_directory to current directory
+    config["working_directory"] = os.getcwd()
+    return config
 
-    # Ensure 'working_directory' exists, default if not
-    if "working_directory" not in config or not config["working_directory"]:
-        config["working_directory"] = os.getcwd()
+def save_config(config_path: str, config: dict) -> None:
+    """Save config dict (with updated working_directory) to config_path."""
+    with open(config_path, 'w') as f:
+        json.dump(config, f, indent=4)
 
-    # Check if model configuration is needed
-    if not config.get("model_name") or not config.get("llm_provider"):
-        print("LLM provider or model not configured.")
-        questions = [
-            inquirer.List('provider_key',
-                message="Select LLM Provider",
-                choices=list(AVAILABLE_MODELS.keys()) # User selects "openai", "mistral", etc.
-            )
-        ]
-        selected_provider_key = inquirer.prompt(questions)['provider_key']
-        clear_terminal()
-        
-        # Ensure API key exists for the selected provider
-        ensure_api_key(selected_provider_key) # Uses "openai", "mistral", etc.
-        
-        questions = [
-            inquirer.List('model_name_full',
-                message=f"Select {selected_provider_key} Model",
-                choices=AVAILABLE_MODELS[selected_provider_key] # Shows "openai/gpt-3.5-turbo", etc.
-            )
-        ]
-        selected_model_name_full = inquirer.prompt(questions)['model_name_full']
+# --- API Key Management ---
+
+def ensure_api_key(provider: str):
+    """Ensure that an API key for the given provider (e.g. "openai") is available (via .env or prompt) and return it. Raise an error if unknown provider."""
+    env_path = os.path.join(os.path.dirname(__file__), '../.env')
+    env_var = API_KEYS.get(provider)
+    if not env_var:
+        raise ValueError(f"Unknown provider: {provider}")
+
+    # Force reload .env (if it exists) so that any new key is picked up.
+    if os.path.exists(env_path):
+        load_dotenv(env_path, override=True)
+
+    api_key = os.getenv(env_var)
+    if not api_key:
+        questions = [ inquirer.Text("api_key", message=f"Enter your {provider.upper()} API key", validate=lambda _, x: len(x.strip()) > 0) ]
+        api_key = inquirer.prompt(questions)["api_key"]
         clear_terminal()
-        
-        config["llm_provider"] = selected_provider_key # Store "openai"
-        config["model_name"] = selected_model_name_full # Store "openai/gpt-4o"
-        
-        with open(config_path, 'w') as f:
-            json.dump(config, f, indent=4)
-        print(f"Configuration saved: Provider '{selected_provider_key}', Model '{selected_model_name_full}'")
+        # Save (or update) the key in .env
+        lines = []
+        if os.path.exists(env_path):
+             with open(env_path, 'r') as f:
+                 lines = f.readlines()
+        key_line = f"{env_var}={api_key}\n"
+        key_exists = False
+        for (i, line) in enumerate(lines):
+             if (line.strip().startswith(f"{env_var}=") or line.strip().startswith(f"#{env_var}=")):
+                 lines[i] = key_line
+                 key_exists = True
+                 break
+        if not key_exists:
+             lines.append(key_line)
+        with open(env_path, 'w') as f:
+             f.writelines(lines)
+        # Reload .env (override) so that the new key is available.
+        load_dotenv(env_path, override=True)
+    
 
-    else:
-        # Config exists, ensure API key for the stored provider
-        # llm_provider should already be the provider part, e.g., "openai"
-        # If old config only had model_name, try to parse provider from it
-        provider_to_check = config.get("llm_provider")
-        if not provider_to_check and config.get("model_name"):
-            provider_to_check = get_provider_from_model_name(config["model_name"])
-            # Optionally, update config if llm_provider was missing
-            if provider_to_check != config.get("llm_provider"): # Check if it's different or was None
-                config["llm_provider"] = provider_to_check
-                with open(config_path, 'w') as f:
-                    json.dump(config, f, indent=4)
+# --- Model / Provider Management ---
 
-        if provider_to_check:
-             ensure_api_key(provider_to_check)
-        else:
-            # This case should ideally be handled by the initial setup logic
-            print("Warning: Could not determine LLM provider from config to ensure API key.")
+def prompt_model_selection() -> tuple:
+    """Prompt the user (via inquirer) to select a provider (from AVAILABLE_MODELS) and then a model (from that provider's list). Return (provider, model_name_full)."""
+    questions = [ inquirer.List("provider_key", message="Select LLM Provider", choices=list(AVAILABLE_MODELS.keys())) ]
+    selected_provider_key = inquirer.prompt(questions)["provider_key"]
+    clear_terminal()
+    # (Ensure API key for the selected provider.)
+    ensure_api_key(selected_provider_key)
+    questions = [ inquirer.List("model_name_full", message=f"Select {selected_provider_key} Model", choices=AVAILABLE_MODELS[selected_provider_key]) ]
+    selected_model_name_full = inquirer.prompt(questions)["model_name_full"]
+    clear_terminal()
+    return (selected_provider_key, selected_model_name_full)
 
+def update_model_config(config_path: str, provider_key: str, model_name_full: str) -> None:
+    """Update config (at config_path) so that "llm_provider" is provider_key and "model_name" is model_name_full. (Also update "working_directory" to os.getcwd() if missing.)"""
+    config = load_config(config_path)
+    config["llm_provider"] = provider_key
+    config["model_name"] = model_name_full
+    if "working_directory" not in config or not config["working_directory"]:
+         config["working_directory"] = os.getcwd()
+    save_config(config_path, config)
 
-    # Create config file if it was created from scratch without example
-    if not os.path.exists(config_path):
-        with open(config_path, 'w') as f:
-            json.dump(config, f, indent=4)
-            
-    return config_path
+# --- CLI Commands ---
 
-@click.group(invoke_without_command=True)
-@click.option('--task', '-t', help='The task to automate')
-@click.option('--max-iter', '-m', default=10, help='Maximum number of iterations for the task')
-@click.option('--change-model', is_flag=True, help='Change the LLM provider and model')
+@click.group(invoke_without_command=True, help="TaskAutomator – Your AI Task Automation Tool\n\nThis tool helps automate tasks using AI. You can run tasks directly or use various commands to manage settings and tools.")
+@click.option("--task", "-t", help="The task to automate (e.g., 'create a python script that sorts files by date')")
+@click.option("--max-iter", "-m", default=10, help="Maximum number of iterations for the task (default: 10)")
+@click.option("--change-model", is_flag=True, help="Change the LLM provider and model before running the task")
 @click.pass_context
 def cli(ctx, task, max_iter, change_model):
-    """TaskAutomator - Your AI Task Automation Tool"""
-    # Ensure config exists and has required fields
-    config_path = ensure_config_exists()
-    clear_terminal()
-    
-    # If change-model flag is set, update the model
-    if change_model:
-        with open(config_path, 'r') as f:
-            config = json.load(f)
-        
-        print("Current configuration: Provider: {}, Model: {}".format(config.get("llm_provider"), config.get("model_name")))
-        questions = [
-            inquirer.List('provider_key',
-                message="Select LLM Provider",
-                choices=list(AVAILABLE_MODELS.keys()) # User selects "openai", "mistral", etc.
-            )
-        ]
-        selected_provider_key = inquirer.prompt(questions)['provider_key']
-        clear_terminal()
-        
-        # Ensure API key exists for the selected provider
-        ensure_api_key(selected_provider_key)
-        
-        questions = [
-            inquirer.List('model_name_full',
-                message=f"Select {selected_provider_key} Model",
-                choices=AVAILABLE_MODELS[selected_provider_key] # Shows "openai/gpt-3.5-turbo", etc.
-            )
-        ]
-        selected_model_name_full = inquirer.prompt(questions)['model_name_full']
-        clear_terminal()
-        
-        config["llm_provider"] = selected_provider_key # Store "openai"
-        config["model_name"] = selected_model_name_full # Store "openai/gpt-4o"
-        
-        # Ensure working_directory is preserved or set
-        if "working_directory" not in config or not config["working_directory"]:
-            config["working_directory"] = os.getcwd()
-
-        with open(config_path, 'w') as f:
-            json.dump(config, f, indent=4)
-        
-        click.echo(f"Model changed to {selected_model_name_full}")
-        return
+    """TaskAutomator – Your AI Task Automation Tool
     
-    # Ensure API key for the configured model before running OpenCopilot
-    # This is a bit redundant if ensure_config_exists already did it, but good for safety
-    with open(config_path, 'r') as f:
-        config = json.load(f)
+    This tool helps automate tasks using AI. You can:
+    - Run tasks directly with --task
+    - Change AI models with --change-model
+    - Manage API keys with set-api-key and set-serp-key
+    - List available tools and models
+    """
+    config_path = os.path.join(os.path.dirname(__file__), '../config.json')
+    config = load_config(config_path)
+    save_config(config_path, config)
     
+    clear_terminal()
+
+    if change_model or not config.get("model_name"):
+         (provider, model) = prompt_model_selection()
+         update_model_config(config_path, provider, model)
+         click.echo(f"Model changed to {model}")
+         if not change_model:  # Only return if this was triggered by missing model
+              return
+
+    # Ensure API key for the configured model (or derive provider from model_name if missing) before running OpenCopilot.
     current_provider = config.get("llm_provider")
-    if not current_provider and config.get("model_name"): # If llm_provider is missing, try to derive it
-        current_provider = get_provider_from_model_name(config["model_name"])
-    
+    if not current_provider and config.get("model_name"):
+         current_provider = get_provider_from_model_name(config["model_name"])
     if current_provider:
-        ensure_api_key(current_provider)
+         ensure_api_key(current_provider)
     else:
-        click.echo("Error: LLM provider not configured. Please run with --change-model to set it up.", err=True)
-        return
+         click.echo("Error: LLM provider not configured. Please run with --change-model to set it up.", err=True)
+         return
 
     copilot = OpenCopilot()
     if ctx.invoked_subcommand is None:
-        if task:
-            copilot.run_task(user_prompt=task, max_iter=max_iter)
-        else:
-            copilot.run()
+         if task:
+             copilot.run_task(user_prompt=task, max_iter=max_iter)
+         else:
+             copilot.run()
 
-@cli.command('list-tools')
+@cli.command("list-tools", help="List all available automation tools and their descriptions")
 def list_tools():
-    """List all available automation tools"""
+    """List all available automation tools and their descriptions.
+    
+    This command shows all tools that can be used by the AI to automate tasks,
+    including what each tool does and what arguments it accepts.
+    """
     tools = OpenCopilot.list_available_tools()
     click.echo("Available Tools:")
     for tool in tools:
-        click.echo(f"- {tool['name']}: {tool['summary']}")
-        if tool.get('arguments'):
-            click.echo(f"    Arguments: {tool['arguments']}")
+         click.echo(f"- {tool['name']}: {tool['summary']}")
+         if tool.get("arguments"):
+             click.echo(f"    Arguments: {tool['arguments']}")
 
-@cli.command('list-models')
+@cli.command("list-models", help="List all available LLM providers and their models")
 def list_models():
-    """List all available LLM providers and their models (litellm compatible)"""
+    """List all available LLM providers and their models.
+    
+    Shows all supported AI models that can be used for task automation,
+    organized by provider (OpenAI, Mistral, Groq, Anthropic).
+    """
     click.echo("Available LLM Providers and Models (litellm compatible):")
-    for provider_key, model_list in AVAILABLE_MODELS.items():
-        click.echo(f"\n{provider_key.upper()}:") # provider_key is "openai", "mistral", etc.
-        for model_name_full in model_list: # model_name_full is "openai/gpt-4o", etc.
-            click.echo(f"  - {model_name_full}")
+    for (provider_key, model_list) in AVAILABLE_MODELS.items():
+         click.echo(f"\n{provider_key.upper()}:")
+         for model_name_full in model_list:
+             click.echo(f"  - {model_name_full}")
 
-@cli.command('set-api-key')
-@click.option('--provider', '-p', type=click.Choice(list(AVAILABLE_MODELS.keys())), 
-              help='The LLM provider to set API key for')
-@click.option('--key', '-k', help='The API key to set (if not provided, will prompt for it)')
+@cli.command("set-api-key", help="Set or update API key for an LLM provider")
+@click.option("--provider", "-p", type=click.Choice(list(AVAILABLE_MODELS.keys())), help="The LLM provider to set API key for (e.g., openai, mistral, groq, anthropic)")
+@click.option("--key", "-k", help="The API key to set (if not provided, will prompt for it securely)")
 def set_api_key(provider, key):
-    """Set or update API key for a specific LLM provider"""
-    if not provider:
-        # If no provider specified, ask user to choose
-        questions = [
-            inquirer.List('provider_key',
-                message="Select LLM Provider to update API key",
-                choices=list(AVAILABLE_MODELS.keys())
-            )
-        ]
-        provider = inquirer.prompt(questions)['provider_key']
+    """Set or update API key for an LLM provider.
+    
+    This command allows you to set or update the API key for any supported LLM provider.
+    The key will be stored securely in your .env file.
     
-    # Get the environment variable name for this provider
+    Examples:
+        piko set-api-key --provider openai
+        piko set-api-key -p mistral -k your-key-here
+    """
+    if not provider:
+         questions = [ inquirer.List("provider_key", message="Select LLM Provider to update API key", choices=list(AVAILABLE_MODELS.keys())) ]
+         provider = inquirer.prompt(questions)["provider_key"]
     env_var = API_KEYS.get(provider)
     if not env_var:
-        raise ValueError(f"Unknown provider: {provider}")
-    
-    # Get the API key (either from command line or prompt)
+         raise ValueError(f"Unknown provider: {provider}")
     if not key:
-        questions = [
-            inquirer.Text('api_key',
-                message=f"Enter your {provider.upper()} API key",
-                validate=lambda _, x: len(x.strip()) > 0
-            )
-        ]
-        key = inquirer.prompt(questions)['api_key']
-    
-    # Get the path to .env file
+         questions = [ inquirer.Text("api_key", message=f"Enter your {provider.upper()} API key", validate=lambda _, x: len(x.strip()) > 0) ]
+         key = inquirer.prompt(questions)["api_key"]
     env_path = os.path.join(os.path.dirname(__file__), '../.env')
-    
-    # Read existing .env file if it exists
     lines = []
     if os.path.exists(env_path):
-        with open(env_path, 'r') as f:
-            lines = f.readlines()
-    
-    # Update or add the API key
+         with open(env_path, 'r') as f:
+             lines = f.readlines()
     key_line = f"{env_var}={key}\n"
     key_exists = False
+    for (i, line) in enumerate(lines):
+         if (line.strip().startswith(f"{env_var}=") or line.strip().startswith(f"#{env_var}=")):
+             lines[i] = key_line
+             key_exists = True
+             break
+    if not key_exists:
+         lines.append(key_line)
+    with open(env_path, 'w') as f:
+         f.writelines(lines)
+    click.echo(f"API key for {provider.upper()} has been updated successfully in {env_path}")
+
+@cli.command("set-serp-key", help="Set or update the SERP API key for web search functionality")
+@click.option("--key", "-k", help="The SERP API key to set (if not provided, will prompt for it securely)")
+def set_serp_key(key):
+    """Set or update the SERP API key used for web search functionality.
     
-    for i, line in enumerate(lines):
-        if line.strip().startswith(f"{env_var}=") or line.strip().startswith(f"#{env_var}="):
-            lines[i] = key_line
-            key_exists = True
-            break
+    This command sets the API key used for web search operations when DuckDuckGo
+    search is not available. The key will be stored securely in your .env file.
     
+    Examples:
+        piko set-serp-key
+        piko set-serp-key -k your-key-here
+    """
+    if not key:
+         questions = [ inquirer.Text("api_key", message="Enter your SERP API key", validate=lambda _, x: len(x.strip()) > 0) ]
+         key = inquirer.prompt(questions)["api_key"]
+    env_path = os.path.join(os.path.dirname(__file__), '../.env')
+    lines = []
+    if os.path.exists(env_path):
+         with open(env_path, 'r') as f:
+             lines = f.readlines()
+    key_line = f"SERP_API_KEY={key}\n"
+    key_exists = False
+    for (i, line) in enumerate(lines):
+         if (line.strip().startswith("SERP_API_KEY=") or line.strip().startswith("#SERP_API_KEY=")):
+             lines[i] = key_line
+             key_exists = True
+             break
     if not key_exists:
-        lines.append(key_line)
-    
-    # Write back to .env file
+         lines.append(key_line)
     with open(env_path, 'w') as f:
-        f.writelines(lines)
-    
-    click.echo(f"API key for {provider.upper()} has been updated successfully in {env_path}")
+         f.writelines(lines)
+    click.echo(f"SERP API key has been updated successfully in {env_path}")
 
 if __name__ == '__main__':
     cli() 

llm_interface/llm.py

@@ -7,11 +7,26 @@ import os
 import sys
 import json
 import litellm # Added import for litellm
+import logging
+from datetime import datetime
 
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')))
 
 from Utils.ter_interface import TerminalInterface
 
+# Set up logging
+log_dir = os.path.join(os.path.dirname(__file__), '../../logs')
+os.makedirs(log_dir, exist_ok=True)
+log_file = os.path.join(log_dir, 'llm_responses.log')
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(message)s',
+    handlers=[
+        logging.FileHandler(log_file)  # Only log to file, removed StreamHandler
+    ]
+)
 
 # Load environment variables from .env file
 load_dotenv()
@@ -30,6 +45,15 @@ class LiteLLMInterface:
     def __init__(self):
         self.terminal = TerminalInterface()
         self.model_name = self.load_config()
+        logging.info(f"\n{'='*50}\nNew Session - Using model: {self.model_name}\n{'='*50}")
+
+    def log_response(self, response_content):
+        """Log only the LLM response in a readable format."""
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        log_entry = f"\n{'='*50}\nTimestamp: {timestamp}\nModel: {self.model_name}\n\n"
+        log_entry += f"Response:\n{response_content}\n"
+        log_entry += f"{'='*50}\n"
+        logging.info(log_entry)
 
     def load_config(self):
         config_path = os.path.join(os.path.dirname(__file__), '../../config.json')
@@ -60,10 +84,12 @@ class LiteLLMInterface:
                     response_content += content
             
             self.terminal.flush_markdown()
+            # Log only the response after successful completion
+            self.log_response(response_content)
             return response_content
         except Exception as e:
-            # litellm maps exceptions to OpenAI exceptions.
-            # The executor should catch these and handle them.
+            # Log the error
+            logging.error(f"\n{'='*50}\nError occurred:\nModel: {self.model_name}\nError: {str(e)}\n{'='*50}")
             print(f"An error occurred during the API call: {e}")
             self.terminal.flush_markdown() # Ensure terminal is flushed even on error
             raise # Re-raise the exception to be caught by the executor

{pikoai-0.1.16.dist-info → pikoai-0.1.18.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: PikoAi
-Version: 0.1.16
+Version: 0.1.18
 Summary: An AI-powered task automation tool
 Home-page: https://github.com/nihaaaar22/OS-Assistant
 Author: Nihar S

{pikoai-0.1.16.dist-info → pikoai-0.1.18.dist-info}/RECORD

@@ -1,9 +1,9 @@
 OpenCopilot.py,sha256=kPTs0-ly84h4dM7AmBlK4uwst5Sj2AM6UAlE3okkD8U,12157
-cli.py,sha256=qDWlKyOUkhDHSBxQ2wjy1e6KGRsDxE7pjxUZKjQFC3k,13959
+cli.py,sha256=2UvmH74pcBFFezI0WHNyWTHMYasIM5NGnrUX6wsdveM,12945
 Agents/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 Agents/Executor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-Agents/Executor/executor.py,sha256=OAqrDdjO62Gy4UKEVF0yYYJjkvVJo_mQy_OMJSibl58,8255
-Agents/Executor/prompts.py,sha256=dWmUiAloNN_NKwWgaG_IqHDhCX1O-VhwbadUL-16gZw,2903
+Agents/Executor/executor.py,sha256=BzSYkT4aPW1yDLSNXNr9WEWZEcs1becEYYlop-eB8s8,6999
+Agents/Executor/prompts.py,sha256=pGY4uXNGYiw_TnTUsRjrVsWc9CV657q3916eui0oulU,2688
 Env/__init__.py,sha256=KLe7UcNV5L395SxhMwbYGyu7KPrSNaoV_9QJo3mLop0,196
 Env/base_env.py,sha256=K4PoWwPXn3pKeu7_-JOlUuyNbyYQ9itMhQybFOm-3K4,1563
 Env/base_executor.py,sha256=awTwJ44CKWV4JO2KUHfHDX0p1Ujw55hlaL5oNYTEW9M,893
@@ -18,18 +18,18 @@ Tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 Tools/file_task.py,sha256=VUhWq_G-SWvGahQo8PG7TOpElHUW3BGLUabrTdJS89o,12151
 Tools/system_details.py,sha256=RScVnhTpUOlNG0g5bGnwmtNr5nSJzOec8HJSFpbicds,2651
 Tools/tool_dir.json,sha256=RTawcanxIkJaUys6Y3yftXAT5uxMH0xPZYTtD1ilJl0,3119
-Tools/tool_manager.py,sha256=86qwREw5an12dweIGCS1NNgINwHikyTxUpbjPWoLbt0,4118
+Tools/tool_manager.py,sha256=xrjrGmLI1hXqCfALBRBuMFCgXwc7pb5TpKMBH_4tCzs,4106
 Tools/userinp.py,sha256=SK69fMEdUvNQor9V3BVckeDMJcq71g6H6EHPmNfsZD4,834
 Tools/web_loader.py,sha256=_oP48uwveTaCKU7G5ju2zsJGTcZd1ScXTKOvHDFtZJU,4564
-Tools/web_search.py,sha256=4EGq1VZqfDgG-_yXTd4_Ha1iEUcR-szdlgRV7oFPru4,1259
+Tools/web_search.py,sha256=12_VhwJGXmn3oUNhTbQ5ENFG964t9DWkfCz3UtlxrbM,2261
 Utils/__init__.py,sha256=oukU0ufroPRd8_N8d2xiFes9CTxSaw4NA6p2nS1kkSg,16
 Utils/executor_utils.py,sha256=WwK3TKgw_hG_crg7ijRaqfidYnnNXYbbs37vKZRYK-0,491
-Utils/ter_interface.py,sha256=2k32kVxcxVBewXnjSGSPbx25ZLhKindEMtL6wSC_wL4,3829
+Utils/ter_interface.py,sha256=8Oe5818MAYC21SfUxtfnV9HQFcQ49z8Q030jjPqNP_g,3889
 llm_interface/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-llm_interface/llm.py,sha256=tI_KDOW14QLWowA7bB3GPe2qjlk0sjS5fBavs9XD1fo,5185
-pikoai-0.1.16.dist-info/licenses/LICENSE,sha256=cELUVOboOAderKFp8bdtcM5VyJi61YH1oDbRhOuoQZw,1067
-pikoai-0.1.16.dist-info/METADATA,sha256=218K7ggeGVwRAv1lubC0_rATLNEpMWUVBP8eFYmenb0,2962
-pikoai-0.1.16.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-pikoai-0.1.16.dist-info/entry_points.txt,sha256=xjZnheDymNDnQ0o84R0jZKEITrhNbzQWN-AhqfA_d6s,50
-pikoai-0.1.16.dist-info/top_level.txt,sha256=hWzBNE7UQsuNcENIOksGcJED08k3ZGRRn2X5jnStICU,53
-pikoai-0.1.16.dist-info/RECORD,,
+llm_interface/llm.py,sha256=TEKKgRfZGtk_UgBfqqkdWKHr4NS3jOOk9di0S3RA7c0,6216
+pikoai-0.1.18.dist-info/licenses/LICENSE,sha256=cELUVOboOAderKFp8bdtcM5VyJi61YH1oDbRhOuoQZw,1067
+pikoai-0.1.18.dist-info/METADATA,sha256=BM1IQxmZ4D7p07EfgLAuUGhpLvAT3W9oPw9WTV8Lee0,2962
+pikoai-0.1.18.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+pikoai-0.1.18.dist-info/entry_points.txt,sha256=xjZnheDymNDnQ0o84R0jZKEITrhNbzQWN-AhqfA_d6s,50
+pikoai-0.1.18.dist-info/top_level.txt,sha256=hWzBNE7UQsuNcENIOksGcJED08k3ZGRRn2X5jnStICU,53
+pikoai-0.1.18.dist-info/RECORD,,