File size: 670 Bytes
2029b44 a51fed4 2029b44 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
from naptha_sdk.utils import get_logger
logger = get_logger(__name__)
def run(inputs, worker_nodes = None, orchestrator_node = None, flow_run = None, cfg: dict = None):
from litellm import completion
prompt = inputs
messages = [
{"role": "system", "content": "You are a helpful AI assistant."},
{"role": "user", "content": prompt},
]
response = completion(
model="ollama/llama3.1:70b",
messages=messages,
temperature=0,
max_tokens=1000,
api_base="http://localhost:11434",
)
response = response.choices[0].message["content"]
logger.info(f"Response: {response}")
return response |