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.
235 lines
7.6 KiB
235 lines
7.6 KiB
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}")
|