File size: 850 Bytes
86f9adc 07d3694 86f9adc 07d3694 86f9adc 07d3694 86f9adc 07d3694 86f9adc 07d3694 86f9adc 07d3694 9892dcf 07d3694 9892dcf 07d3694 86f9adc 07d3694 86f9adc | 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 | import os
from dotenv import load_dotenv
load_dotenv()
from ollama import chat
from groq import Groq
ENV = os.getenv("ENV")
if ENV == "production":
inference_client = Groq(api_key=os.environ.get("GROQ_INFERENCE_KEY"))
def llm(user_prompt: str, system_prompt: str = None) -> str:
messages = []
if system_prompt:
messages.append({"role": "system", "content": system_prompt})
messages.append({"role": "user", "content": user_prompt})
if ENV == "local":
response = chat(
"llama3.2:3b",
messages=messages
)
return response["message"]["content"]
elif ENV == "production":
response = inference_client.chat.completions.create(
messages=messages,
model="llama-3.3-70b-versatile"
)
return response.choices[0].message.content
|