import requests from typing import Any, Dict, List, Optional from langchain.llms.base import LLM from langchain.agents import initialize_agent, AgentType, Tool from pydantic import Field, BaseModel import os from dotenv import load_dotenv from datetime import datetime import wikipedia from asteval import Interpreter # For a safer calculator import logging # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Load environment variables from .env file if present load_dotenv() # Constants for LM Studio API LM_STUDIO_API_URL = os.getenv("LM_STUDIO_API_URL", "http://192.168.0.104:1234/v1/chat/completions") MODEL_NAME = os.getenv("LM_STUDIO_MODEL", "lmstudio-community/Meta-Llama-3.1-8B-Instruct-GGUF/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf") CONTENT_TYPE = "application/json" class LMStudioLLM(LLM): """ Custom LangChain LLM to interface with LM Studio API. """ api_url: str = Field(default=LM_STUDIO_API_URL, description="The API endpoint for LM Studio.") model: str = Field(default=MODEL_NAME, description="The model path/name.") temperature: float = Field(default=0.7, description="Sampling temperature.") max_tokens: Optional[int] = Field(default=4096, description="Maximum number of tokens to generate.") streaming: bool = Field(default=False, alias="stream", description="Whether to use streaming responses.") class Config: populate_by_name = True @property def _llm_type(self) -> str: return "lmstudio" @property def identifying_params(self) -> Dict[str, Any]: return { "api_url": self.api_url, "model": self.model, "temperature": self.temperature, "max_tokens": self.max_tokens, "stream": self.streaming, } def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str: """ Generate a response from the LM Studio model. Args: prompt (str): The input prompt. stop (Optional[List[str]]): Stop sequences. Returns: str: The generated response. """ headers = { "Content-Type": CONTENT_TYPE, } payload = { "model": self.model, "messages": [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": prompt}, ], "temperature": self.temperature, "max_tokens": self.max_tokens if self.max_tokens is not None else -1, "stream": self.streaming, # Uses alias 'stream' } logger.info(f"Payload: {payload}") try: response = requests.post( self.api_url, headers=headers, json=payload, timeout=60, stream=self.streaming ) response.raise_for_status() logger.info(f"Response content: {response.text}") except requests.RequestException as e: logger.error(f"Failed to connect to LM Studio API: {e}") raise RuntimeError(f"Failed to connect to LM Studio API: {e}") if self.streaming: return self._handle_stream(response) else: try: response_json = response.json() choices = response_json.get("choices", []) if not choices: raise ValueError("No choices found in the response.") # Extract the first response's content content = choices[0].get("message", {}).get("content", "") return content.strip() except (ValueError, KeyError) as e: logger.error(f"Invalid response format: {e}") raise RuntimeError(f"Invalid response format: {e}") def _handle_stream(self, response: requests.Response) -> str: """ Process streaming responses from the LM Studio API. Args: response (requests.Response): The streaming response object. Returns: str: The concatenated content from the stream. """ content = "" try: for line in response.iter_lines(): if line: decoded_line = line.decode('utf-8') if decoded_line.startswith("data: "): data = decoded_line[6:] if data == "[DONE]": break try: json_data = requests.utils.json.loads(data) choices = json_data.get("choices", []) for chunk in choices: delta = chunk.get("delta", {}) piece = delta.get("content", "") content += piece except requests.utils.json.JSONDecodeError: continue return content.strip() except Exception as e: logger.error(f"Error processing streaming response: {e}") raise RuntimeError(f"Error processing streaming response: {e}") def calculator(input: str) -> str: """ A simple calculator tool that safely evaluates mathematical expressions. Args: input (str): The mathematical expression to evaluate. Returns: str: The result of the evaluation. """ try: aeval = Interpreter() result = aeval(input) return str(result) except Exception as e: logger.error(f"Calculator error: {e}") return f"Error evaluating expression: {e}" def wikipedia_search(query: str) -> str: """ Search Wikipedia for a given query and return a summary. Args: query (str): The search query. Returns: str: A summary from Wikipedia. """ try: summary = wikipedia.summary(query, sentences=2) return summary except wikipedia.exceptions.DisambiguationError as e: return f"Disambiguation error. Options include: {e.options}" except wikipedia.exceptions.PageError: return "No page found for the query." except Exception as e: logger.error(f"Wikipedia search error: {e}") return f"An error occurred: {e}" def current_time(_input: str) -> str: """ Returns the current system time. Args: _input (str): Ignored input. Returns: str: The current time as a string. """ now = datetime.now() return now.strftime("%Y-%m-%d %H:%M:%S") # Initialize the custom LLM llm = LMStudioLLM() # Define Tools tools = [ Tool( name="Calculator", func=calculator, description="Useful for performing mathematical calculations. Input should be a valid mathematical expression.", ), Tool( name="WikipediaSearch", func=wikipedia_search, description="Useful for fetching summaries from Wikipedia. Input should be a search query.", ), Tool( name="CurrentTime", func=current_time, description="Returns the current system time.", ), ] # Initialize the Agent agent = initialize_agent( tools=tools, llm=llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True, handle_parsing_errors=True, ) # Example Usage if __name__ == "__main__": user_input = "What is the capital of France and what is 15 multiplied by 3?" try: response = agent({"input": user_input}) print(response["output"]) except Exception as e: logger.error(f"Agent invocation error: {e}") print(f"An error occurred: {e}")