File size: 1,575 Bytes
6e8a7d4 e207857 102a9b5 4260a74 102a9b5 4260a74 102a9b5 4260a74 102a9b5 4260a74 102a9b5 4260a74 6e8a7d4 4260a74 e207857 6e8a7d4 28e2398 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
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() |