매뉴얼 등록 분리

manual.py

@@ -1,10 +1,12 @@
 import os
+from threading import Thread
 
 import torch
 from sentence_transformers import SentenceTransformer
 import chromadb
+import json
 from chromadb.utils import embedding_functions
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 from fastapi import FastAPI
 
 # 2. 벡터 DB 설정
@@ -99,12 +101,46 @@ def query_and_summarize(job: str, query: str, top_k: int = 3):
     print("\n\n\n\n\n")
     return content
 
+def query_and_summarize_stream(job: str, query: str):
+    # 1. 문서 검색 (기존과 동일)
+    results = collection.query(query_texts=[query], n_results=1, where={"dept": job})
+    top_doc = results['documents'][0][0]
+
+    model, tokenizer = get_qwen_model()
+
+    # 2. 메시지 구성
+    messages = [
+        {"role": "system", "content": "당신은 회사 재무/회계 업무 전문 어시스턴트입니다."},
+        {"role": "user", "content": f"다음 문서를 참고하세요:\n{top_doc}\n\n질문: {query}"}
+    ]
+
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+    # 3. 스트리머 설정
+    # skip_prompt=True를 해야 입력한 질문이 다시 나오지 않습니다.
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    # 4. 별도 스레드에서 생성 실행 (비동기 처리를 위함)
+    generation_kwargs = dict(model_inputs, streamer=streamer, max_new_tokens=500)
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    # 5. 제너레이터 함수 정의
+    def generate():
+        for new_text in streamer:
+            if new_text:
+                # 클라이언트가 JSON으로 받길 원한다면 형식을 맞춰줍니다.
+                yield json.dumps({"kind": "text", "text": new_text}) + "\n"
+
+    return generate
+
 app = FastAPI()
 
 @app.get("/")
 def question(query: str) :
     user_query = query
-    answer = query_and_summarize(job="FI", query=user_query)
+    answer = query_and_summarize_stream(job="FI", query=user_query)
     return {"answer": answer}
 
 # 예시 사용