diff --git a/swarms/utils/llm.py b/swarms/utils/llm.py index e69de29b..b6d6d0c6 100644 --- a/swarms/utils/llm.py +++ b/swarms/utils/llm.py @@ -0,0 +1,63 @@ +from typing import Optional +import os +import logging + +from langchain import PromptTemplate, HuggingFaceHub, ChatOpenAI, LLMChain + +# Configure logging +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +class LLM: + def __init__(self, + openai_api_key: Optional[str] = None, + hf_repo_id: Optional[str] = None, + hf_api_token: Optional[str] = None, + model_kwargs: Optional[dict] = None): + + # Check if keys are in the environment variables + openai_api_key = openai_api_key or os.getenv('OPENAI_API_KEY') + hf_api_token = hf_api_token or os.getenv('HUGGINGFACEHUB_API_TOKEN') + + self.openai_api_key = openai_api_key + self.hf_repo_id = hf_repo_id + self.hf_api_token = hf_api_token + self.model_kwargs = model_kwargs if model_kwargs else {} + + # If the HuggingFace API token is provided, set it in environment variables + if self.hf_api_token: + os.environ["HUGGINGFACEHUB_API_TOKEN"] = self.hf_api_token + + # Create the LLM object based on the provided keys + try: + if self.hf_repo_id and self.hf_api_token: + self.llm = HuggingFaceHub(repo_id=self.hf_repo_id, model_kwargs=self.model_kwargs) + elif self.openai_api_key: + self.llm = ChatOpenAI(api_key=self.openai_api_key, model_kwargs=self.model_kwargs) + else: + raise ValueError("Please provide either OpenAI API key or both HuggingFace repository ID and API token.") + except Exception as e: + logger.error("Failed to initialize LLM: %s", e) + raise + + def run(self, prompt: str) -> str: + template = """Question: {question} + Answer: Let's think step by step.""" + + try: + prompt_template = PromptTemplate(template=template, input_variables=["question"]) + llm_chain = LLMChain(prompt=prompt_template, llm=self.llm) + return llm_chain.run({"question": prompt}) + except Exception as e: + logger.error("Failed to generate response: %s", e) + raise + +# example +llm_instance = LLM(openai_api_key="your_openai_key") +result = llm_instance.run("Who won the FIFA World Cup in 1998?") +print(result) + +# using HuggingFaceHub +llm_instance = LLM(hf_repo_id="google/flan-t5-xl", hf_api_token="your_hf_api_token") +result = llm_instance.run("Who won the FIFA World Cup in 1998?") +print(result)