70 lines
1.9 KiB
Python
70 lines
1.9 KiB
Python
import asyncio
|
|
import json
|
|
import toml
|
|
from loguru import logger
|
|
import nats
|
|
from providers import get_provider
|
|
|
|
# Загрузка конфигурации
|
|
config = toml.load("config.toml")
|
|
|
|
NATS_URL = config["nats"]["url"]
|
|
INCOMING_TOPIC = config["nats"]["incoming_topic"]
|
|
RESPONSE_TOPIC = config["nats"]["response_topic"]
|
|
|
|
provider = get_provider(
|
|
config["llm"]["provider"],
|
|
api_key=config["llm"]["api_key"],
|
|
model=config["llm"]["model"],
|
|
system_prompt=config["llm"].get("system_prompt", "")
|
|
)
|
|
|
|
|
|
async def message_handler(msg):
|
|
logger.info(f"Received message on {msg.subject}: {msg.data}")
|
|
|
|
try:
|
|
data = json.loads(msg.data.decode())
|
|
request_id = data.get("request_id")
|
|
user_id = data.get("user_id")
|
|
text = data.get("text")
|
|
if not text:
|
|
logger.warning("Empty text field in message")
|
|
return
|
|
|
|
response = await provider.chat([{"role": "user", "content": text}])
|
|
|
|
reply = {
|
|
"request_id": request_id,
|
|
"user_id": user_id,
|
|
"response": response,
|
|
"service": "LLM-Service"
|
|
}
|
|
|
|
await nc.publish(RESPONSE_TOPIC, json.dumps(reply).encode())
|
|
logger.success(f"Ответ опубликован в {RESPONSE_TOPIC} для request_id: {request_id}")
|
|
|
|
except Exception as e:
|
|
logger.exception("Error handling message")
|
|
|
|
|
|
async def main():
|
|
logger.info("Connecting to NATS...")
|
|
global nc
|
|
nc = await nats.connect(servers=[NATS_URL])
|
|
await nc.subscribe(INCOMING_TOPIC, cb=message_handler)
|
|
await nc.flush()
|
|
logger.info(f"Subscribed to {TOPIC_IN}")
|
|
|
|
# Просто держим процесс активным
|
|
try:
|
|
logger.info("LLM service is running. Waiting for messages...")
|
|
await asyncio.Event().wait()
|
|
except KeyboardInterrupt:
|
|
logger.info("Shutting down LLM service...")
|
|
finally:
|
|
await nc.drain()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|