Axcel1/icd_10_coding_assistant
1
1from fastapi import FastAPI, HTTPException, Query2from fastapi.middleware.cors import CORSMiddleware3from fastapi.responses import JSONResponse4from pydantic import BaseModel5from typing import List, Optional, Dict, Any6import time7import logging8import pprint9 10# Import your existing neural searcher and the new multi-collection system11# from neural_searcher import NeuralSearcher12from chapter_retrieval_system_v2 import MultiCollectionChapterRetrieval13 14# Configure logging15logging.basicConfig(level=logging.INFO)16logger = logging.getLogger(__name__)17 18app = FastAPI(19 title="ICD-10 Multi-Collection Search API",20 description="Advanced ICD-10 code search with intelligent chapter detection",21 version="2.0.0"22)23 24# Add CORS middleware for web frontend integration25app.add_middleware(26 CORSMiddleware,27 allow_origins=["*"], # Configure this properly for production28 allow_credentials=True,29 allow_methods=["*"],30 allow_headers=["*"],31)32 33# Initialize systems34try:35 # Initialize the multi-collection chapter retrieval system36 chapter_retriever = MultiCollectionChapterRetrieval()37 38 # Keep your original neural searcher for backward compatibility39 # You might not need this if switching fully to multi-collection approach40 # neural_searcher = NeuralSearcher(collection_name="icd10_codes_chapter_3")41 42 logger.info("Successfully initialized search systems")43except Exception as e:44 logger.error(f"Failed to initialize search systems: {e}")45 chapter_retriever = None46 # neural_searcher = None47 48# Pydantic models for request/response validation49class SearchRequest(BaseModel):50 query: str51 limit: Optional[int] = 1052 score_threshold: Optional[float] = 0.353 search_mode: Optional[str] = "smart" # "smart", "all_chapters", "specific_chapters"54 target_chapters: Optional[List[str]] = None55 detailed_analysis: Optional[bool] = False56 chapters_per_sentence: Optional[int] = 2 # NEW: How many chapters to search per sentence57 58 59 60class ChapterInfo(BaseModel):61 chapter_id: str62 collection_name: str63 relevance_score: float64 description: str65 match_count: int66 avg_score: float67 max_score: float68 69class SearchResult(BaseModel):70 code: str71 title: str72 description: Optional[str] = None73 score: float74 chapter_id: Optional[str] = None75 collection: str76 source_sentence: Optional[str] = None # NEW: Track which sentence generated this result77 sentence_key: Optional[str] = None # NEW: Track sentence identifier78 79class SentenceResults(BaseModel):80 sentence_text: str81 sentence_key: str82 results: List[SearchResult]83 total_results: int84 85class SearchResponse(BaseModel):86 query: str87 total_results: int88 search_time: float89 search_mode: str90 relevant_chapters: List[ChapterInfo]91 results: List[SearchResult] # Keep for backward compatibility92 sentence_results: Optional[List[SentenceResults]] = None # NEW: Results grouped by sentence93 94 95class ChapterAnalysisResponse(BaseModel):96 query: str97 analysis_time: float98 chapters: List[ChapterInfo]99 100# Health check endpoint101@app.get("/health")102def health_check():103 """Health check endpoint"""104 if chapter_retriever is None:105 raise HTTPException(status_code=503, detail="Search system not initialized")106 return {"status": "healthy", "timestamp": time.time()}107 108# Chapter analysis endpoint109@app.get("/api/analyze-chapters", response_model=ChapterAnalysisResponse)110def analyze_chapters(111 q: str = Query(..., description="Diagnostic query string"),112 detailed: bool = Query(False, description="Include detailed chapter statistics")113):114 """115 Analyze which ICD-10 chapters are most relevant for a diagnostic query116 """117 if not chapter_retriever:118 raise HTTPException(status_code=503, detail="Chapter retrieval system not available")119 120 if not q or not q.strip():121 raise HTTPException(status_code=400, detail="Query parameter 'q' is required")122 123 try:124 start_time = time.time()125 126 # Perform chapter analysis127 analysis = chapter_retriever.analyze_chapters_parallel(128 q.strip(),129 sample_size_per_chapter=15,130 score_threshold=0.2131 )132 133 analysis_time = time.time() - start_time134 135 # Convert to response format136 chapters = []137 for chapter_id, stats in analysis.items():138 if stats['relevance_score'] > 0.05: # Filter very low relevance139 chapter_info = ChapterInfo(140 chapter_id=chapter_id,141 collection_name=stats['collection_name'],142 relevance_score=stats['relevance_score'],143 description=chapter_retriever.chapter_info.get(chapter_id, "Unknown chapter"),144 match_count=stats['match_count'],145 avg_score=stats['avg_score'],146 max_score=stats['max_score']147 )148 chapters.append(chapter_info)149 150 return ChapterAnalysisResponse(151 query=q,152 analysis_time=analysis_time,153 chapters=chapters154 )155 156 except Exception as e:157 logger.error(f"Error in chapter analysis: {e}")158 raise HTTPException(status_code=500, detail=f"Chapter analysis failed: {str(e)}")159 160# Smart search endpoint (main search functionality)161@app.post("/api/search", response_model=SearchResponse)162def search_smart(request: SearchRequest):163 """164 Advanced search with intelligent chapter detection and targeted searching165 """166 return _perform_search(request)167 168@app.get("/api/search", response_model=SearchResponse)169def search_smart_get(170 q: str = Query(..., description="Diagnostic query string"),171 limit: int = Query(10, ge=1, le=100, description="Maximum number of results"),172 score_threshold: float = Query(0.3, ge=0.0, le=1.0, description="Minimum similarity score"),173 search_mode: str = Query("smart", description="Search mode: smart, all_chapters, specific_chapters"),174 target_chapters: Optional[str] = Query(None, description="Comma-separated list of target chapters (for specific_chapters mode)"),175 detailed_analysis: bool = Query(False, description="Include detailed chapter analysis"),176 chapters_per_sentence: int = Query(2, ge=1, le=5, description="Number of chapters to search per sentence") # NEW177):178 """179 Advanced search with intelligent chapter detection (GET version)180 """181 # Parse target_chapters if provided182 parsed_chapters = None183 if target_chapters:184 parsed_chapters = [ch.strip() for ch in target_chapters.split(",") if ch.strip()]185 186 request = SearchRequest(187 query=q,188 limit=limit,189 score_threshold=score_threshold,190 search_mode=search_mode,191 target_chapters=parsed_chapters,192 detailed_analysis=detailed_analysis,193 chapters_per_sentence=chapters_per_sentence # NEW194 )195 196 return _perform_search(request)197 198def _perform_search(request: SearchRequest) -> SearchResponse:199 """Internal search logic - UPDATED to return top responses for each sentence"""200 if not chapter_retriever:201 raise HTTPException(status_code=503, detail="Search system not available")202 203 if not request.query or not request.query.strip():204 raise HTTPException(status_code=400, detail="Query is required")205 206 try:207 start_time = time.time()208 query = request.query.strip()209 210 # Initialize response data211 relevant_chapters = []212 results = []213 sentence_results = [] # NEW: For sentence-based results214 215 if request.search_mode == "smart":216 # Smart search: auto-identify chapters then search them sentence by sentence217 logger.info(f"Performing sentence-based smart search for: '{query}'")218 219 # First, analyze chapters if detailed analysis is requested220 if request.detailed_analysis:221 analysis = chapter_retriever.analyze_chapters_parallel(query)222 for chapter_id, stats in analysis.items():223 if stats['relevance_score'] > 0.1:224 chapter_info = ChapterInfo(225 chapter_id=chapter_id,226 collection_name=stats['collection_name'],227 relevance_score=stats['relevance_score'],228 description=chapter_retriever.chapter_info.get(chapter_id, "Unknown"),229 match_count=stats['match_count'],230 avg_score=stats['avg_score'],231 max_score=stats['max_score']232 )233 relevant_chapters.append(chapter_info)234 235 # Perform sentence-based targeted search236 search_results = chapter_retriever.search_targeted_chapters(237 query, 238 target_chapters=request.target_chapters,239 results_per_sentence=request.limit, # Use full limit per sentence240 chapters_per_sentence=request.chapters_per_sentence241 )242 243 # NEW: Process results by sentence instead of flattening244 sentence_result_map = {} # Track results by sentence245 all_results = [] # Keep flattened results for backward compatibility246 247 # Group results by sentence248 for chapter_id, chapter_data in search_results.items():249 for sentence_key, sentence_data in chapter_data.items():250 sentence_text = sentence_data['text']251 252 # Initialize sentence entry if not exists253 if sentence_key not in sentence_result_map:254 sentence_result_map[sentence_key] = {255 'text': sentence_text,256 'results': []257 }258 259 # Add results for this sentence260 for result in sentence_data['results']:261 # Create enriched result with metadata262 enriched_result = {263 **result,264 'chapter_id': chapter_id,265 'source_sentence': sentence_text,266 'sentence_key': sentence_key267 }268 269 # Add to sentence-specific results270 sentence_result_map[sentence_key]['results'].append(enriched_result)271 272 # Add to flattened results for backward compatibility273 all_results.append(enriched_result)274 275 # NEW: Create sentence-based result objects276 for sentence_key, sentence_data in sentence_result_map.items():277 # Sort sentence results by score278 sentence_data['results'].sort(key=lambda x: x['score'], reverse=True)279 280 # Apply score threshold and limit per sentence281 filtered_sentence_results = [282 r for r in sentence_data['results'] 283 if r['score'] >= request.score_threshold284 ][:request.limit]285 286 # Convert to SearchResult objects287 sentence_search_results = []288 for result in filtered_sentence_results:289 payload = result['payload']290 search_result = SearchResult(291 code=payload.get('code', 'N/A'),292 title=payload.get('title', 'N/A'),293 description=payload.get('description'),294 score=result['score'],295 chapter_id=result.get('chapter_id'),296 collection=result['collection'],297 source_sentence=result.get('source_sentence'),298 sentence_key=result.get('sentence_key')299 )300 sentence_search_results.append(search_result)301 302 # Create SentenceResults object303 if sentence_search_results: # Only include sentences with results304 sentence_result_obj = SentenceResults(305 sentence_text=sentence_data['text'],306 sentence_key=sentence_key,307 results=sentence_search_results,308 total_results=len(sentence_search_results)309 )310 sentence_results.append(sentence_result_obj)311 312 # Sort sentence results by average score (optional)313 sentence_results.sort(314 key=lambda x: sum(r.score for r in x.results) / len(x.results) if x.results else 0,315 reverse=True316 )317 318 # Process flattened results for backward compatibility319 all_results.sort(key=lambda x: x['score'], reverse=True)320 all_results = all_results[:request.limit]321 322 elif request.search_mode == "all_chapters":323 # Handle other search modes (keeping original logic)324 # You can implement similar sentence-based logic here if needed325 logger.info("All chapters search mode - using original logic")326 # ... implement if needed327 328 elif request.search_mode == "specific_chapters":329 # Handle specific chapters mode330 logger.info("Specific chapters search mode - using original logic")331 # ... implement if needed332 333 else:334 raise HTTPException(status_code=400, detail=f"Unknown search mode: {request.search_mode}")335 336 # Convert flattened results to response format (for backward compatibility)337 for result in all_results:338 if result['score'] >= request.score_threshold:339 payload = result['payload']340 search_result = SearchResult(341 code=payload.get('code', 'N/A'),342 title=payload.get('title', 'N/A'),343 description=payload.get('description'),344 score=result['score'],345 chapter_id=result.get('chapter_id'),346 collection=result['collection'],347 source_sentence=result.get('source_sentence'),348 sentence_key=result.get('sentence_key')349 )350 results.append(search_result)351 352 search_time = time.time() - start_time353 354 logger.info(f"Sentence-based search completed: {len(results)} total results, {len(sentence_results)} sentences in {search_time:.3f}s")355 356 # Debug output357 logger.info(f"Sentence results breakdown:")358 for sent_result in sentence_results:359 logger.info(f" '{sent_result.sentence_text}': {sent_result.total_results} results")360 361 return SearchResponse(362 query=query,363 total_results=len(results),364 search_time=search_time,365 search_mode=request.search_mode,366 relevant_chapters=relevant_chapters,367 results=results, # Flattened results for backward compatibility368 sentence_results=sentence_results # NEW: Results organized by sentence369 )370 371 except Exception as e:372 logger.error(f"Search error: {e}")373 raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}")374 375 376 377 378# Backward compatibility endpoint (your original endpoint)379# @app.get("/api/search/legacy")380# def search_legacy(q: str):381# """382# Legacy search endpoint for backward compatibility383# Uses your original neural searcher384# """385# # if not neural_searcher:386# # raise HTTPException(status_code=503, detail="Legacy search system not available")387 388# if not q or not q.strip():389# raise HTTPException(status_code=400, detail="Query parameter 'q' is required")390 391# try:392# result = neural_searcher.search(text=q.strip())393# return {"result": result}394# except Exception as e:395# logger.error(f"Legacy search error: {e}")396# raise HTTPException(status_code=500, detail=f"Legacy search failed: {str(e)}")397 398# Get available chapters399@app.get("/api/chapters")400def get_available_chapters():401 """402 Get list of available ICD-10 chapters and their descriptions403 """404 if not chapter_retriever:405 raise HTTPException(status_code=503, detail="Chapter system not available")406 407 try:408 chapter_collections = chapter_retriever.get_chapter_collections()409 410 chapters = []411 for chapter_id, collection_name in chapter_collections.items():412 description = chapter_retriever.chapter_info.get(chapter_id, "Unknown chapter")413 chapters.append({414 "chapter_id": chapter_id,415 "collection_name": collection_name,416 "description": description417 })418 419 return {420 "total_chapters": len(chapters),421 "chapters": chapters422 }423 except Exception as e:424 logger.error(f"Error getting chapters: {e}")425 raise HTTPException(status_code=500, detail=f"Failed to get chapters: {str(e)}")426 427# Get search suggestions/autocomplete (optional enhancement)428@app.get("/api/suggest")429def get_search_suggestions(430 q: str = Query(..., min_length=2, description="Partial query for suggestions"),431 limit: int = Query(5, ge=1, le=20, description="Maximum number of suggestions")432):433 """434 Get search suggestions based on partial query435 This is a simple implementation - you might want to enhance this436 """437 # Simple keyword-based suggestions438 # In a real implementation, you might use a more sophisticated approach439 440 common_terms = [441 "chest pain", "shortness of breath", "diabetes", "hypertension", 442 "pneumonia", "fracture", "depression", "anxiety", "fever",443 "headache", "abdominal pain", "nausea", "vomiting", "infection",444 "cancer", "tumor", "heart attack", "stroke", "asthma"445 ]446 447 query_lower = q.lower().strip()448 suggestions = [term for term in common_terms if query_lower in term.lower()]449 450 return {"suggestions": suggestions[:limit]}451 452if __name__ == "__main__":453 import uvicorn454 455 # Run with more configuration options456 uvicorn.run(457 app, 458 host="0.0.0.0", 459 port=8000,460 log_level="info",461 access_log=True462 )