Update helpers/foundation_models.py

helpers/foundation_models.py

@@ -1,10 +1,11 @@
 import os
-from typing import List, Tuple
+from typing import List, Tuple, Dict, Any
 
 import openai
 import streamlit as st
 from langchain.agents import AgentType, initialize_agent, load_tools
 from langchain.llms import OpenAI as l_OpenAI
+import requests
 
 OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
 SERPAPI_API_KEY = os.environ["SERPAPI_API_KEY"]
@@ -43,13 +44,106 @@ def call_chatgpt(query: str, model: str = "gpt-3.5-turbo") -> str:
 
 
 def call_langchain(prompt: str) -> str:
-    llm = l_OpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)
-    tools = load_tools(
+    """
+    Initializes a language model with specific settings, loads additional tools, initializes an agent with these tools,
+    and then runs the agent with a given prompt to produce a text response.
+
+    Args:
+    prompt (str): The input text prompt that the agent will process.
+
+    Returns:
+    str: The text output produced by the agent after processing the input prompt.
+    """
+
+    # Initialize the OpenAI language model with a specified temperature
+    # and the OpenAI API key. It's assumed that `l_OpenAI` is a class or function
+    # that is responsible for setting up the language model with the given parameters.
+    llm = l_OpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)  # Type: ignore
+
+    # Load additional tools needed for the agent. Here, 'serpapi' for search engine results
+    # and 'llm-math' for math capabilities are loaded, along with their respective API keys.
+    # The `load_tools` function is assumed to return a dictionary of initialized tools,
+    # with the tools being ready for use by the agent.
+    tools = load_tools(  # Type: ignore
         ["serpapi", "llm-math"], llm=llm, serpapi_api_key=SERPAPI_API_KEY
     )
-    agent = initialize_agent(
-        tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
+
+    # Initialize the agent with the provided tools, the language model, and specific agent settings.
+    # The agent is set to a ZERO_SHOT_REACT_DESCRIPTION type, which likely defines its behavior
+    # or capabilities, with verbosity enabled for detailed logs.
+    # The `initialize_agent` function presumably returns an instance of an agent configured
+    # with the specified tools and settings, ready to process prompts.
+    agent = initialize_agent(  # Type: ignore
+        tools, llm, agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
     )
-    output = agent.run(prompt)
+
+    # Execute the agent with the given prompt and capture the output.
+    # The `run` method of the agent is assumed to process the prompt and return a string response,
+    # which is then returned by this function to the caller.
+    output: str = agent.run(prompt)
 
     return output
+
+
+def query(payload: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Sends a JSON payload to a predefined API URL and returns the JSON response.
+
+    Args:
+        payload (Dict[str, Any]): The JSON payload to be sent to the API.
+
+    Returns:
+        Dict[str, Any]: The JSON response received from the API.
+    """
+
+    # API endpoint URL
+    API_URL = "https://sks7h7h5qkhoxwxo.us-east-1.aws.endpoints.huggingface.cloud"
+    
+    # Headers to indicate both the request and response formats are JSON
+    headers = {
+        "Accept": "application/json",
+        "Content-Type": "application/json"
+    }
+    
+    # Sending a POST request with the JSON payload and headers
+    response = requests.post(API_URL, headers=headers, json=payload)
+    
+    # Returning the JSON response
+    return response.json()
+
+def llama2_7b_ysa(prompt: str) -> str:
+    """
+    Queries a model and retrieves the generated text based on the given prompt.
+
+    This function sends a prompt to a model (presumably named 'llama2_7b') and extracts
+    the generated text from the model's response. It's tailored for handling responses
+    from a specific API or model query structure where the response is expected to be
+    a list of dictionaries, with at least one dictionary containing a key 'generated_text'.
+
+    Parameters:
+    - prompt (str): The text prompt to send to the model.
+
+    Returns:
+    - str: The generated text response from the model.
+
+    Note:
+    - The function assumes that the 'query' function is previously defined and accessible
+      within the same scope or module. It should send a request to the model and return
+      the response in a structured format.
+    - The 'parameters' dictionary is passed empty but can be customized to include specific
+      request parameters as needed by the model API.
+    """
+
+    # Define the query payload with the prompt and any additional parameters
+    query_payload: Dict[str, Any] = {
+        "inputs": prompt,
+        "parameters": {}
+    }
+
+    # Send the query to the model and store the output response
+    output = query(query_payload)
+
+    # Extract the 'generated_text' from the first item in the response list
+    response: str = output[0]['generated_text']
+
+    return response