|
|
import streamlit as st |
|
|
import pandas as pd |
|
|
from smolagents import CodeAgent, tool |
|
|
from groq import Groq |
|
|
import os |
|
|
|
|
|
class GroqLLM: |
|
|
def __init__(self, model_name="llama-3.1-8B-Instant"): |
|
|
self.client = Groq(api_key=os.environ.get("GROQ_API_KEY")) |
|
|
self.model_name = model_name |
|
|
|
|
|
def __call__(self, prompt: str): |
|
|
completion = self.client.chat.completions.create( |
|
|
model=self.model_name, |
|
|
messages=[{"role": "user", "content": prompt}], |
|
|
temperature=0.7, |
|
|
max_tokens=1024, |
|
|
stream=False, |
|
|
) |
|
|
return completion.choices[0].message.content |
|
|
|
|
|
@tool |
|
|
def execute_code(code_string: str, data: pd.DataFrame) -> str: |
|
|
"""Executes python code and returns results as a string. |
|
|
|
|
|
Args: |
|
|
code_string (str): Python code to execute. |
|
|
data (pd.DataFrame): The dataframe to use in the code |
|
|
Returns: |
|
|
str: The result of executing the code or an error message |
|
|
""" |
|
|
try: |
|
|
local_vars = {"data": data, "pd": pd} |
|
|
exec(code_string, local_vars) |
|
|
if "result" in local_vars: |
|
|
return str(local_vars["result"]) |
|
|
return "Success, but no 'result' variable assigned." |
|
|
except Exception as e: |
|
|
return f"Error: {e}" |
|
|
|
|
|
def main(): |
|
|
st.title("Test") |
|
|
data = pd.DataFrame({"a": [1, 2], "b": [3, 4]}) |
|
|
|
|
|
agent = CodeAgent( |
|
|
tools=[execute_code], |
|
|
model=GroqLLM(), |
|
|
) |
|
|
code = "result = data.sum()" |
|
|
result = agent.run(f"Use the execute_code tool, run this python code ```python\n{code}\n```", data=data) |
|
|
st.write(result) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
main() |
