Update app.py

app.py

@@ -107,26 +107,178 @@ def respond(
         return [{"role": "user", "content": message},
                 {"role": "assistant", "content": f"오류가 발생했습니다: {str(e)}"}]
 
+def analyze_prompt(message):
+    """LLM을 사용하여 프롬프트 분석"""
+    try:
+        analysis_prompt = f"""
+        다음 요청을 분석하여 단백질 설계에 필요한 주요 특성을 추출하세요:
+        요청: {message}
+        
+        다음 항목들을 분석해주세요:
+        1. 주요 기능 (예: 치료, 결합, 촉매 등)
+        2. 목표 환경 (예: 세포막, 수용성, 등)
+        3. 필요한 구조적 특징
+        4. 크기 및 복잡도 요구사항
+        """
+        
+        response = client.chat.completions.create(
+            model="CohereForAI/c4ai-command-r-plus-08-2024",
+            messages=[{"role": "user", "content": analysis_prompt}],
+            temperature=0.7
+        )
+        
+        return response.choices[0].message.content
+    except Exception as e:
+        print(f"프롬프트 분석 중 오류: {str(e)}")
+        return None
+
+def search_protein_data(analysis, dataset):
+    """분석 결과를 바탕으로 데이터셋에서 유사한 구조 검색"""
+    try:
+        # 분석에서 키워드 추출
+        keywords = extract_keywords(analysis)
+        
+        # 유사도 점수 계산
+        scored_entries = []
+        for entry in dataset['train']:
+            score = calculate_similarity(keywords, entry)
+            scored_entries.append((score, entry))
+        
+        # 상위 유사 구조 반환
+        scored_entries.sort(reverse=True)
+        return scored_entries[:3]  # 상위 3개 구조 반환
+    except Exception as e:
+        print(f"데이터 검색 중 오류: {str(e)}")
+        return []
+
+def extract_parameters(analysis, similar_structures):
+    """분석 결과와 유사 구조를 바탕으로 생성 파라미터 결정"""
+    try:
+        # 기본 파라미터 템플릿
+        params = {
+            'sequence_length': 100,
+            'helix_bias': 0.02,
+            'strand_bias': 0.02,
+            'loop_bias': 0.1,
+            'hydrophobic_target_score': 0
+        }
+        
+        # 분석 결과에서 구조적 요구사항 파악
+        if "막 투과" in analysis or "소수성" in analysis:
+            params['hydrophobic_target_score'] = -2
+            params['helix_bias'] = 0.03
+        elif "수용성" in analysis or "가용성" in analysis:
+            params['hydrophobic_target_score'] = 2
+            params['loop_bias'] = 0.15
+            
+        # 유사 구조들의 특성 반영
+        if similar_structures:
+            avg_length = sum(len(s[1]['sequence']) for s in similar_structures) / len(similar_structures)
+            params['sequence_length'] = int(avg_length)
+            
+            # 구조적 특성 분석 및 반영
+            for _, structure in similar_structures:
+                if 'secondary_structure' in structure:
+                    helix_ratio = structure['secondary_structure'].count('H') / len(structure['secondary_structure'])
+                    sheet_ratio = structure['secondary_structure'].count('E') / len(structure['secondary_structure'])
+                    params['helix_bias'] = max(0.01, min(0.05, helix_ratio))
+                    params['strand_bias'] = max(0.01, min(0.05, sheet_ratio))
+        
+        return params
+    except Exception as e:
+        print(f"파라미터 추출 중 오류: {str(e)}")
+        return None
+
 def process_chat(message, history):
     try:
-        # 단백질 생성 관련 키워드 확인
-        if any(keyword in message.lower() for keyword in ['protein', 'generate', '단백질', '생성']):
-            # 데이터셋에서 관련 정보 검색
-            relevant_data = search_protein_data(message)
+        if any(keyword in message.lower() for keyword in ['protein', 'generate', '단백질', '생성', '치료']):
+            # 1. LLM을 사용한 프롬프트 분석
+            analysis = analyze_prompt(message)
+            if not analysis:
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": "요청 분석에 실패했습니다."}
+                ]
             
-            # 생성 파라미터 추출
-            params = extract_parameters(message, relevant_data)
+            # 2. 유사 구조 검색
+            similar_structures = search_protein_data(analysis, ds)
+            if not similar_structures:
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": "적합한 참조 구조를 찾지 못했습니다."}
+                ]
             
-            # 단백질 생성
-            protein_result = generate_protein(params)
+            # 3. 생성 파라미터 결정
+            params = extract_parameters(analysis, similar_structures)
+            if not params:
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": "파라미터 설정에 실패했습니다."}
+                ]
             
-            # 결과 설명 생성
-            explanation = generate_explanation(protein_result, params)
+            # 4. 단백질 생성
+            try:
+                protein_result = protein_diffusion_model(
+                    sequence=None,
+                    seq_len=params['sequence_length'],
+                    helix_bias=params['helix_bias'],
+                    strand_bias=params['strand_bias'],
+                    loop_bias=params['loop_bias'],
+                    secondary_structure=None,
+                    aa_bias=None,
+                    aa_bias_potential=None,
+                    num_steps="25",
+                    noise="normal",
+                    hydrophobic_target_score=str(params['hydrophobic_target_score']),
+                    hydrophobic_potential="2",
+                    contigs=None,
+                    pssm=None,
+                    seq_mask=None,
+                    str_mask=None,
+                    rewrite_pdb=None
+                )
+                
+                output_seq, output_pdb, structure_view, plddt_plot = next(protein_result)
+                
+                # 5. 결과 설명 생성
+                explanation = f"""
+                요청하신 기능에 맞는 단백질을 생성했습니다:
+                
+                분석된 요구사항:
+                {analysis}
+                
+                설계된 구조적 특징:
+                - 길이: {params['sequence_length']} 아미노산
+                - 알파 헬릭스 비율: {params['helix_bias']*100:.1f}%
+                - 베타 시트 비율: {params['strand_bias']*100:.1f}%
+                - 루프 구조 비율: {params['loop_bias']*100:.1f}%
+                - 소수성 점수: {params['hydrophobic_target_score']}
+                
+                참조된 유사 구조: {len(similar_structures)}개
+                
+                생성된 단백질의 3D 구조와 시퀀스를 확인하실 수 있습니다.
+                """
+                
+                # 6. 결과 저장
+                global current_protein_result
+                current_protein_result = {
+                    'sequence': output_seq,
+                    'pdb': output_pdb,
+                    'structure_view': structure_view,
+                    'plddt_plot': plddt_plot,
+                    'params': params
+                }
+                
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": explanation}
+                ]
             
-            return history + [
-                {"role": "user", "content": message},
-                {"role": "assistant", "content": explanation}
-            ]
+            except Exception as e:
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": f"단백질 생성 중 오류가 발생했습니다: {str(e)}"}
+                ]
         else:
             return history + [
                 {"role": "user", "content": message},
@@ -135,27 +287,9 @@ def process_chat(message, history):
     except Exception as e:
         return history + [
             {"role": "user", "content": message},
-            {"role": "assistant", "content": f"오류가 발생했습니다: {str(e)}"}
+            {"role": "assistant", "content": f"처리 중 오류가 발생했습니다: {str(e)}"}
         ]
 
-def search_protein_data(query):
-    # 데이터셋에서 관련 항목 검색
-    relevant_entries = []
-    for entry in ds['train']:
-        if any(keyword in entry['sequence'].lower() for keyword in query.lower().split()):
-            relevant_entries.append(entry)
-    return relevant_entries
-
-def extract_parameters(llm_response, dataset_info):
-    # LLM 응답에서 파라미터 추출
-    params = {
-        'sequence_length': 100,  # 기본값
-        'helix_bias': 0.02,
-        'strand_bias': 0.02,
-        'loop_bias': 0.1,
-        'hydrophobic_target_score': 0
-    }
-    return params
 
 def generate_protein(params):
     # 기존 protein_diffusion_model 함수 호출