Merge pull request #39 from studio-recoding/feat/recommendation

[fix] recommendation main 테스트 완료
studio-recoding · Mar 31, 2024 · 2e75853 · 2e75853
2 parents 55f776a + d75cc06
commit 2e75853
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 15 deletions.
diff --git a/app/database/chroma_db.py b/app/database/chroma_db.py
@@ -57,19 +57,23 @@ async def add_db_data(schedule_data: AddScheduleDTO):
 # 메인페이지 한 줄 추천 기능에 사용하는 함수
 # 유저의 id, 해당 날짜로 필터링
 async def db_recommendation_main(user_data: RecommendationMainRequestDTO):
+    member = user_data.member_id
+    schedule_datetime_start = user_data.schedule_datetime_start
+    schedule_datetime_end = user_data.schedule_datetime_end
     results = schedules.query(
-        user_persona=["hard working"],
+        query_texts=["hard working"],
         n_results=5,
         where={"$and" :
                [
-                   {"member": {"$eq": int(user_data.member_id)}},
+                   {"member": {"$eq": int(member)}},
                    {"datetime_start": {
-                       "$gte": user_data.schedule_datetime_start, # greater than or equal
-                       "$lt": user_data.schedule_datetime_end # less than
+                       "$eq": schedule_datetime_start, # greater than or equal
+                       # "$lt": schedule_datetime_end # less than
                    }}
                ]}
         # where_document={"$contains":"search_string"}  # optional filter
     )
+    return results['documents']
 
 def get_chroma_client():
     return chroma_client
diff --git a/app/prompt/openai_prompt.py b/app/prompt/openai_prompt.py
@@ -6,14 +6,8 @@ class Template:
                     3. The recommendation must be concise, limited to one sentence without any additional commentary.
                     
                     Example:
-                    User schedule:  8:00 AM - 9:00 AM: Gym
-                                    9:30 AM - 12:00 PM: Work meetings
-                                    12:00 PM - 1:00 PM: Lunch break
-                                    1:00 PM - 5:00 PM: Work on project
-                                    5:30 PM - 7:00 PM: Free time
-                                    7:00 PM - 9:00 PM: Dinner with family
-                                    9:30 PM: Free time
-                    AI Recommendation: "Since you've got some free time before dinner, how about taking a short walk in the park to relax and clear your mind?"
+                    User schedule:  [Practice guitar, Calculate accuracy, Study backend development, Run AI models in the lab, Study NEST.JS]
+                    AI Recommendation: "Your day is filled with learning and research. how about taking a short walk in between studies?"
                     
                     User schedule: {schedule}
                     AI Recommendation:

diff --git a/app/routers/recommendation.py b/app/routers/recommendation.py
@@ -25,11 +25,11 @@
 config = configparser.ConfigParser()
 config.read(CONFIG_FILE_PATH)
 
-@router.get("/main", status_code=status.HTTP_200_OK)
+@router.post("/main", status_code=status.HTTP_200_OK)
 async def get_recommendation(user_data: RecommendationMainRequestDTO) -> ChatResponse:
 
     # 모델
-    chat_model = ChatOpenAI(temperature=0,  # 창의성 (0.0 ~ 2.0)
+    chat_model = ChatOpenAI(temperature=0.5,  # 창의성 (0.0 ~ 2.0)
                             max_tokens=2048,  # 최대 토큰수
                             model_name='gpt-3.5-turbo-1106',  # 모델명
                             openai_api_key=OPENAI_API_KEY  # API 키
@@ -38,10 +38,12 @@ async def get_recommendation(user_data: RecommendationMainRequestDTO) -> ChatRes
     # vectordb에서 유저의 정보를 가져온다.
     schedule = await vectordb.db_recommendation_main(user_data)
 
+    print(schedule)
+
     # 템플릿
     recommendation_template = openai_prompt.Template.recommendation_template
 
     prompt = PromptTemplate.from_template(recommendation_template)
-    result = await chat_model.predict(prompt.format(output_language="Korean", schedule=schedule))
+    result = chat_model.predict(prompt.format(output_language="Korean", schedule=schedule))
     print(result)
     return ChatResponse(ness=result)