From 808caa85c0ae9ba78ecba91f4f66171d6695a457 Mon Sep 17 00:00:00 2001 From: backurs <59131604+backurs@users.noreply.github.com> Date: Wed, 22 Apr 2026 19:55:20 -0400 Subject: [PATCH 1/2] Adding dynamic retrieval pipeline: 3rd attempt (#48) * dynamic retriever with native function calling * fix example yaml file * copilot comments * copilot comments * copilot comments * refactoring and addressing comments --- config_dynamic.yaml.example | 163 ++++++++++++ dynamic_retriever.py | 307 +++++++++++++++++++++++ pubmed/scripts/answer_questions_ref.py | 273 ++++++++++++++++++++ pubmed/scripts/evaluate_answers_ragas.py | 219 ++++++++++++++++ pubmed/scripts/readme.md | 38 +++ utils/cosmos_retriever.py | 132 ++-------- utils/fulltext.py | 69 +++++ utils/ranker.py | 57 +++++ 8 files changed, 1142 insertions(+), 116 deletions(-) create mode 100644 config_dynamic.yaml.example create mode 100644 dynamic_retriever.py create mode 100644 pubmed/scripts/answer_questions_ref.py create mode 100644 pubmed/scripts/evaluate_answers_ragas.py create mode 100644 pubmed/scripts/readme.md create mode 100644 utils/fulltext.py create mode 100644 utils/ranker.py diff --git a/config_dynamic.yaml.example b/config_dynamic.yaml.example new file mode 100644 index 0000000..38bb959 --- /dev/null +++ b/config_dynamic.yaml.example @@ -0,0 +1,163 @@ +# Paths (relative to working directory) +paths: + questions_path: "" + output_root: "" + +prune_k: 20 + +# Azure OpenAI / LLM settings +llm: + llm_endpoint: "" + api_version: "" + llm_model: "" + temperature: 0.0 + max_completion_tokens: 2048 + max_retries: 2 + premium_max_concurrency: 8 + prompt_cache_size: 4096 + use_rbac_auth: true + token_scope: "" + llm_api_key: "" + context_limit: 270000 + +# Embedding settings +embedding: + embed_endpoint: "" + api_version: "" + embed_model: "" + embed_dimensions: 1024 + embed_cache_size: 4096 + use_rbac_auth: false + token_scope: "" + embed_api_key: "" + +# Semantic ranker +ranker: + use_ranker: true + rerank_multiplier: 8 + k_ranker: 10 + region: "" + account_name: "" + read_token_from_path: false + access_token_path: "" + tenant_id: "" + token_scope: "" + register_account_path: "" + url_suffix: "" + batch_size: 16 + max_retries: 10 + +# Cosmos DB settings +cosmos: + uri: "" + key: "" + database_name: "" + use_rbac_auth: false + + # Upload settings + cosmos_account_name: "" + cosmos_resource_group: "" + azure_subscription_id: "" + embedding_batch_size: 20 + + vector_embedding_policy_json: | + { + "vectorEmbeddings": [ + { + "path": "/embedding", + "dataType": "float32", + "dimensions": 1024, + "distanceFunction": "cosine" + } + ] + } + + # Configure each source/container independently + sources: + - id: "source_1" + container_name: "container_1" + partition_key_path: "/pk" + embedding_field: "embedding" + documents_root: "" + embedding_text_fields: + - title + - summary + - content + retrieval: + search_k: 10 + fulltext_search_k: 10 + fulltext_fields: + - title + indexing_policy_json: | + { + "indexingMode": "consistent", + "automatic": true, + "includedPaths": [ + { "path": "/*" } + ], + "excludedPaths": [ + { "path": "/\"_etag\"/?" }, + { "path": "/embedding/*" } + ], + "fullTextIndexes": [ + { "path": "/title" } + ], + "vectorIndexes": [ + { + "path": "/embedding", + "type": "diskANN", + "quantizationByteSize": 192, + "indexingSearchListSize": 100 + } + ] + } + full_text_policy_json: | + { + "defaultLanguage": "en-US", + "fullTextPaths": [ + { "path": "/title", "language": "en-US" } + ] + } + + - id: "source_2" + container_name: "container_2" + partition_key_path: "/pk" + embedding_field: "embedding" + documents_root: "" + embedding_text_fields: + - text + retrieval: + search_k: 25 + fulltext_search_k: 0 + fulltext_fields: + - text + indexing_policy_json: | + { + "indexingMode": "consistent", + "automatic": true, + "includedPaths": [ + { "path": "/*" } + ], + "excludedPaths": [ + { "path": "/\"_etag\"/?" }, + { "path": "/embedding/*" } + ], + "fullTextIndexes": [ + { "path": "/text" } + ], + "vectorIndexes": [ + { + "path": "/embedding", + "type": "diskANN", + "quantizationByteSize": 192, + "indexingSearchListSize": 100 + } + ] + } + full_text_policy_json: | + { + "defaultLanguage": "en-US", + "fullTextPaths": [ + { "path": "/text", "language": "en-US" } + ] + } diff --git a/dynamic_retriever.py b/dynamic_retriever.py new file mode 100644 index 0000000..da7bcd5 --- /dev/null +++ b/dynamic_retriever.py @@ -0,0 +1,307 @@ +"""Standard scaffold with native function calling over Cosmos DB (vector + fulltext + ranker).""" +from tqdm import tqdm +import argparse, asyncio, json, re, time, yaml +from pathlib import Path +import httpx, openai as oai +import tiktoken +from openai import AsyncAzureOpenAI +from azure.identity import AzureCliCredential, get_bearer_token_provider +from azure.identity.aio import AzureCliCredential as AsyncAzureCliCredential +from azure.cosmos.aio import CosmosClient +from utils.fulltext import fulltext_search +from utils.ranker import rerank_documents + +_enc = tiktoken.get_encoding("o200k_base") +def count_tokens(msgs): + return sum(4 + len(_enc.encode(m["content"] if isinstance(m, dict) and "content" in m and isinstance(m["content"], str) else json.dumps(m) if isinstance(m, dict) else str(m))) for m in msgs) + 2 + +# --- Search helpers --- +async def embed(text): + r = await embed_client.embeddings.create(input=[text], model=embed_cfg["embed_model"]) + dim = embed_cfg.get("embed_dimensions", 1536) + raw = [float(x) for x in r.data[0].embedding] + if len(raw) >= dim: + return raw[:dim] + return raw + [0.0] * (dim - len(raw)) + +_SAFE_FIELD_RE = re.compile(r'^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)*$') + +async def vec_search(container, emb, top_k, ef): + if not _SAFE_FIELD_RE.match(ef): + raise ValueError(f"Invalid embedding field name: {ef!r}") + sql = f"SELECT TOP @k c, VectorDistance(c.{ef}, @emb) AS score FROM c ORDER BY VectorDistance(c.{ef}, @emb)" + return [item.get("c", item) async for item in container.query_items(query=sql, parameters=[{"name":"@k","value":top_k},{"name":"@emb","value":emb}])] + +async def rerank(query, docs, top_k): + if not USE_RANKER or not docs: return docs[:top_k] + indices = await rerank_documents(_r_http, _r_url, _r_tok, query, docs, top_k, _r_bs, _r_mr) + if indices is None: + return docs[:top_k] + return [docs[i] for i in indices] + +def fmt(doc): + ex = {"_rid","_self","_etag","_attachments","_ts","_score","e"} | _all_embed + return "\n".join(f"{k}: {v}" for k,v in doc.items() if k not in ex and v) + +async def do_search(query, containers): + emb = await embed(query) + tasks = [] + for sid, ret in _source_cfg.items(): + if sid in containers: + tasks.append(vec_search(containers[sid], emb, ret["search_k"]*RERANK_MUL, _source_embed[sid])) + for sid, fields in _source_ft.items(): + if sid in containers: + tasks.append(fulltext_search(containers[sid], fields, query, _source_cfg[sid]["fulltext_search_k"]*RERANK_MUL)) + results = await asyncio.gather(*tasks) + seen, all_d = set(), [] + for dl in results: + for d in dl: + did = d.get("id","") + if did not in seen: seen.add(did); all_d.append(d) + total_k = sum(r["search_k"]+r["fulltext_search_k"] for r in _source_cfg.values()) + texts = [fmt(d) for d in all_d] + ranked_texts = await rerank(query, texts, total_k) + # Map ranked texts back to source docs by index + text_to_idx = {id(t): i for i, t in enumerate(texts)} + ranked_indices = [text_to_idx[id(t)] for t in ranked_texts] + return json.dumps([{"docid": all_d[i].get("id", ""), "snippet": ranked_texts[j][:2000]} for j, i in enumerate(ranked_indices)]) + +async def do_get_doc(docid, containers): + for container_id, c in containers.items(): + try: + async for item in c.query_items(query="SELECT * FROM c WHERE c.id=@id", parameters=[{"name":"@id","value":docid}]): + return json.dumps({"docid": docid, "text": fmt(item)}) + except asyncio.CancelledError: + raise + except Exception as e: + print(f" [get_document] Cosmos query failed for container={container_id}, docid={docid}: {e}") + continue + return json.dumps({"error": f"Not found: {docid}"}) + +async def do_prune(docids, containers, doc_cache): + parts = [] + for did in docids[:PRUNE_K]: + if did in doc_cache: + parts.append(f'\n{doc_cache[did]}\n') + continue + found = False + for container_id, c in containers.items(): + if found: + break + try: + async for item in c.query_items(query="SELECT * FROM c WHERE c.id=@id", parameters=[{"name":"@id","value":did}]): + text = fmt(item) + doc_cache[did] = text + parts.append(f'\n{text}\n') + found = True + break + except asyncio.CancelledError: + raise + except Exception as e: + print(f" [prune] Cosmos query failed for container={container_id}, docid={did}: {e}") + continue + return "Pruned context (only these documents remain):\n\n" + "\n\n".join(parts) + +TOOLS = [ + {"type":"function","function":{"name":"search","description":"Search knowledge base. Returns top results with docid and snippet.","parameters":{"type":"object","properties":{"query":{"type":"string"}},"required":["query"]}}}, + {"type":"function","function":{"name":"get_document","description":"Get full document by docid.","parameters":{"type":"object","properties":{"docid":{"type":"string"}},"required":["docid"]}}}, + {"type":"function","function":{"name":"prune","description":"Keep only the specified most relevant document IDs and discard all others from context. Use when context is large to free up space for more searches.","parameters":{"type":"object","properties":{"docids":{"type":"array","items":{"type":"string"},"description":"List of document IDs to keep"}},"required":["docids"]}}}, +] + +async def process_question(q_obj, containers): + t0 = time.perf_counter() + query = q_obj["question_text"] + qid = q_obj.get("question_id", "") + print(f"\n{'='*60}\n[{qid}]: {query}\n{'='*60}") + msgs = [{"role": "user", "content": QUERY_TEMPLATE.format(question=query)}] + tc = {"search": 0, "get_document": 0, "prune": 0} + doc_cache = {} + initial_msg = msgs[0] + retries = 0 + for iteration in range(50): + try: + r = await llm.chat.completions.create(model=llm_cfg["llm_model"], messages=msgs, tools=TOOLS, tool_choice="auto", temperature=llm_cfg.get("temperature", 0), max_completion_tokens=llm_cfg["max_completion_tokens"]) + retries = 0 + except (oai.BadRequestError, oai.RateLimitError, oai.APIStatusError) as e: + retries += 1; print(f" LLM error ({retries}/{MAX_RETRIES}): {e}") + if retries >= MAX_RETRIES: + elapsed = round(time.perf_counter() - t0, 2) + print(f" Max retries ({MAX_RETRIES}) exceeded, returning partial result") + return {"question_id": qid, "query": query, "answer": "", "ground_truth": q_obj.get("answer",""), + "model": llm_cfg["llm_model"], "rounds": iteration+1, "elapsed_seconds": elapsed, + "tool_calls": tc, "error": f"Max retries exceeded: {e}"} + await asyncio.sleep(min(5*2**retries, 300)); continue + m = r.choices[0].message + msgs.append(m.model_dump(exclude_none=True)) + if not m.tool_calls: + answer = m.content or "" + print(f" Answer: {answer[:200]}...") + elapsed = round(time.perf_counter() - t0, 2) + print(f" Elapsed: {elapsed}s") + return {"question_id": qid, "query": query, "answer": answer, "ground_truth": q_obj.get("answer",""), + "model": llm_cfg["llm_model"], "rounds": iteration+1, "elapsed_seconds": elapsed, + "tool_calls": tc} + for t in m.tool_calls: + tc[t.function.name] = tc.get(t.function.name, 0) + 1 + # Enforce: prune must be the sole tool call in a turn + call_names = [t.function.name for t in m.tool_calls] + if "prune" in call_names and len(call_names) > 1: + print(f" [warn] prune mixed with other calls; returning error for non-prune calls") + for t in m.tool_calls: + if t.function.name != "prune": + msgs.append({"role": "tool", "tool_call_id": t.id, + "content": json.dumps({"error": "prune must be the only tool call in a turn; re-issue this call separately."})}) + # Execute only the prune call + prune_call = next(t for t in m.tool_calls if t.function.name == "prune") + try: + a = json.loads(prune_call.function.arguments) + except (json.JSONDecodeError, TypeError) as e: + msgs.append({"role": "tool", "tool_call_id": prune_call.id, + "content": json.dumps({"error": f"Malformed tool arguments: {e}"})}) + continue + out = await do_prune(a["docids"], containers, doc_cache) + msgs.clear() + msgs.append(initial_msg) + msgs.append({"role": "assistant", "content": "I'll prune the context to focus on the most relevant documents."}) + msgs.append({"role": "user", "content": out}) + print(f" [prune] Kept {len(a['docids'])} docs, context reset") + continue + async def _exec(t): + try: + a = json.loads(t.function.arguments) + except (json.JSONDecodeError, TypeError) as e: + return t, json.dumps({"error": f"Malformed tool arguments: {e}"}), False + if t.function.name == "search": + out = await do_search(a["query"], containers) + try: + for h in json.loads(out): + if h.get("snippet"): doc_cache[h["docid"]] = h["snippet"] + except asyncio.CancelledError: raise + except Exception as e: + print(f" [search] failed to parse/cache search results: {e}") + elif t.function.name == "get_document": + out = await do_get_doc(a["docid"], containers) + try: + d = json.loads(out) + if d.get("text"): doc_cache[d["docid"]] = d["text"] + except asyncio.CancelledError: raise + except Exception as e: + print(f" [get_document] failed to parse/cache doc result for {a.get('docid','?')}: {e}") + elif t.function.name == "prune": + out = await do_prune(a["docids"], containers, doc_cache) + return t, out, True # signal prune + else: + out = json.dumps({"error": f"Unknown tool: {t.function.name}"}) + return t, out, False + print(f" Executing {len(m.tool_calls)} tool calls in parallel...") + results = await asyncio.gather(*(_exec(t) for t in m.tool_calls)) + pruned = False + for t, out, is_prune in results: + if is_prune: + msgs.clear() + msgs.append(initial_msg) + msgs.append({"role": "assistant", "content": "I'll prune the context to focus on the most relevant documents."}) + msgs.append({"role": "user", "content": out}) + try: + prune_args = json.loads(t.function.arguments) + print(f" [prune] Kept {len(prune_args['docids'])} docs, context reset") + except (json.JSONDecodeError, TypeError, KeyError): + print(f" [prune] context reset") + pruned = True + break + if pruned: + continue + for t, out, _ in results: + msgs.append({"role": "tool", "tool_call_id": t.id, "content": out}) + try: + a = json.loads(t.function.arguments) + print(f" [{t.function.name}] {list(a.values())[0][:80] if isinstance(list(a.values())[0], str) else '...'}") + except (json.JSONDecodeError, TypeError): + print(f" [{t.function.name}] (malformed args)") + token_est = count_tokens(msgs) + msgs.append({"role": "user", "content": f"Token usage: {token_est} / {CONTEXT_LIMIT}"}) + print(f" Token usage: {token_est} / {CONTEXT_LIMIT}") + elapsed = round(time.perf_counter() - t0, 2) + return {"question_id": qid, "query": query, "answer": "", "ground_truth": q_obj.get("answer",""), + "model": llm_cfg["llm_model"], "rounds": 50, "elapsed_seconds": elapsed, "tool_calls": tc} + +async def main(): + questions = json.loads(Path(cfg["paths"]["questions_path"]).read_text()) + if args.max_questions is not None: + questions = questions[:args.max_questions] + use_rbac_auth = cosmos_cfg.get("use_rbac_auth", False) + credential = AsyncAzureCliCredential() if use_rbac_auth else None + cosmos = CosmosClient(cosmos_cfg["uri"], credential=credential or cosmos_cfg["key"]) + db = cosmos.get_database_client(cosmos_cfg["database_name"]) + containers = {s["id"]: db.get_container_client(s["container_name"]) for s in sources} + + try: + results = [] + for q in tqdm(questions): + results.append(await process_question(q, containers)) + + out = Path(cfg["paths"]["output_root"]) / "standard" / f"results_{time.strftime('%Y%m%d_%H%M%S')}.json" + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(json.dumps(results, indent=2)) + print(f"\nSaved {len(results)} results to {out}") + finally: + await cosmos.close() + if credential is not None: + await credential.close() + await llm.close() + await embed_client.close() + if USE_RANKER: + await _r_http.aclose() + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--config", default="config_dynamic.yaml") + parser.add_argument("--max-questions", type=int, default=None, help="Only answer the first N questions") + args = parser.parse_args() + + cfg = yaml.safe_load(Path(args.config).read_text()) + llm_cfg, embed_cfg, cosmos_cfg = cfg["llm"], cfg["embedding"], cfg["cosmos"] + sources = cosmos_cfg["sources"] + _source_cfg = {s["id"]: s["retrieval"] for s in sources} + _source_embed = {s["id"]: s["embedding_field"] for s in sources} + _source_ft = {s["id"]: s["retrieval"]["fulltext_fields"] for s in sources} + _all_embed = set(_source_embed.values()) + MAX_RETRIES, RERANK_MUL = int(llm_cfg["max_retries"]), cfg["ranker"]["rerank_multiplier"] + PRUNE_K = cfg.get("prune_k", 20) + CONTEXT_LIMIT = llm_cfg.get("context_limit", 270000) + + # Clients + tp = get_bearer_token_provider(AzureCliCredential(), llm_cfg["token_scope"]) if llm_cfg["use_rbac_auth"] else None + llm = AsyncAzureOpenAI(api_version=llm_cfg["api_version"], azure_endpoint=llm_cfg["llm_endpoint"], + **({"azure_ad_token_provider": tp} if tp else {"api_key": llm_cfg["llm_api_key"]})) + embed_tp = get_bearer_token_provider(AzureCliCredential(), embed_cfg["token_scope"]) if embed_cfg.get("use_rbac_auth") else None + embed_client = AsyncAzureOpenAI(api_version=embed_cfg["api_version"], azure_endpoint=embed_cfg["embed_endpoint"], + **({"azure_ad_token_provider": embed_tp} if embed_tp else {"api_key": embed_cfg["embed_api_key"]})) + + # Ranker + rcfg = cfg["ranker"] + USE_RANKER = rcfg["use_ranker"] + if USE_RANKER: + _r_url = f"https://{rcfg['account_name']}.{rcfg['region']}.{rcfg['url_suffix']}" + _r_bs, _r_mr = rcfg["batch_size"], rcfg["max_retries"] + if rcfg["read_token_from_path"]: + _r_tok = Path(rcfg["access_token_path"]).read_text().strip() + else: + _ranker_tenant = str(rcfg.get("tenant_id") or "").strip() + _ranker_cred = AzureCliCredential(tenant_id=_ranker_tenant) if _ranker_tenant else AzureCliCredential() + _r_tok = _ranker_cred.get_token(rcfg["token_scope"]).token + _r_hdr = {"Authorization": f"Bearer {_r_tok}", "Content-Type": "application/json"} + _r_http = httpx.AsyncClient(timeout=120) + + QUERY_TEMPLATE = """You are a deep research agent. Answer the question by using the search, get_document, and prune tools. Search multiple times with diverse queries. Do not give up early. + +Available tools: +- search(query): Search the knowledge base. Returns top results with docid and snippet. +- get_document(docid): Get full document text by docid. +- prune(docids): Keep only the specified documents (up to """ + str(PRUNE_K) + """) and discard the rest from context. Use this when context is getting large to focus on the most relevant documents. + +Question: {question} + +Format: Explanation: ... Exact Answer: ... Confidence: N%""" + + asyncio.run(main()) diff --git a/pubmed/scripts/answer_questions_ref.py b/pubmed/scripts/answer_questions_ref.py new file mode 100644 index 0000000..492e78c --- /dev/null +++ b/pubmed/scripts/answer_questions_ref.py @@ -0,0 +1,273 @@ +"""Answer questions: vector search → expand via references → pick closest refs. + +For each question: +1. Vector-search top-k documents (default 10) +2. For each hit, get its references (pmid→pmcid), fetch their embeddings from Cosmos +3. Pick the ref_k references closest to the query vector (default 2) +4. Combine: up to k + k*ref_k documents (e.g. 10 + 20 = 30) +5. Answer with LLM + +Usage: + python answer_questions_ref.py --config config.pubmed.yaml \ + --questions complex_questions.json --output answers_refs.json \ + --top 10 --ref-k 2 +""" + +import argparse, csv, json, logging, sys, time +import numpy as np +import yaml +from azure.cosmos import CosmosClient +from azure.identity import AzureCliCredential, get_bearer_token_provider +from openai import AzureOpenAI +from tqdm import tqdm + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +log = logging.getLogger(__name__) +for n in ("azure.core.pipeline.policies.http_logging_policy", "azure.identity", "azure.cosmos"): + logging.getLogger(n).setLevel(logging.WARNING) + + +def load_pmid_to_pmcid(csv_path: str) -> dict[str, str]: + mapping = {} + with open(csv_path, newline="", encoding="utf-8") as f: + for row in csv.DictReader(f, quotechar='"'): + pmcid = (row.get("PMCID") or "").strip() + pmid = (row.get("PMID") or "").strip() + if pmcid and pmid: + mapping[pmid] = pmcid + return mapping + + +def get_embedding(text: str, embed_client: AzureOpenAI, model: str, dimensions: int, max_input_chars: int = 8000) -> list[float]: + return embed_client.embeddings.create( + input=text[:max_input_chars], model=model, dimensions=dimensions + ).data[0].embedding + + +def vector_search(container, query_vector: list[float], top: int) -> list[dict]: + sql = """ + SELECT TOP @top + c.id, c.pmcid, c.title, c.journal_title, c.abstract, c.pub_year, + c.doi, c.full_text, c.embedding, c.references, + VectorDistance(c.embedding, @embedding) AS similarity + FROM c ORDER BY VectorDistance(c.embedding, @embedding) + """ + return list(container.query_items( + query=sql, + parameters=[{"name": "@embedding", "value": query_vector}, {"name": "@top", "value": top}], + enable_cross_partition_query=True, + )) + + +def fetch_docs_by_pmcids(container, pmcids: list[str]) -> list[dict]: + if not pmcids: + return [] + placeholders = ", ".join(f"@p{i}" for i in range(len(pmcids))) + sql = f""" + SELECT c.id, c.pmcid, c.title, c.journal_title, c.abstract, c.pub_year, + c.doi, c.full_text, c.embedding + FROM c WHERE c.pmcid IN ({placeholders}) + """ + params = [{"name": f"@p{i}", "value": p} for i, p in enumerate(pmcids)] + return list(container.query_items(query=sql, parameters=params, enable_cross_partition_query=True)) + + +def refs_to_pmcids(refs: list[dict], pmid_map: dict[str, str]) -> list[str]: + result = [] + for r in refs: + pmcid = (r.get("pmcid") or "").strip() + if not pmcid: + pmid = (r.get("pmid") or "").strip() + pmcid = pmid_map.get(pmid, "") + if pmcid: + result.append(pmcid) + return result + + +def closest_select(docs: list[dict], query_vec: list[float], k: int) -> list[dict]: + """Select k docs closest to query_vec by cosine similarity.""" + if len(docs) <= k: + return docs + q = np.array(query_vec, dtype=np.float64) + q_norm = np.linalg.norm(q) + if q_norm < 1e-12: + return docs[:k] + q = q / q_norm + scored = [] + for doc in docs: + emb = doc.get("embedding") + if emb and isinstance(emb, list) and len(emb) > 0: + v = np.array(emb, dtype=np.float64) + v_norm = np.linalg.norm(v) + sim = float(np.dot(q, v) / max(v_norm, 1e-12)) + scored.append((sim, doc)) + scored.sort(key=lambda x: x[0], reverse=True) + return [doc for _, doc in scored[:k]] + + +def build_context(docs: list[dict], max_chars: int = 230000) -> str: + parts, total = [], 0 + for i, doc in enumerate(docs, 1): + title = doc.get("title", "(no title)") + pmcid = doc.get("pmcid", "?") + abstract = doc.get("abstract", "") + full_text = doc.get("full_text", "") + sim = doc.get("similarity") + sim_str = f" (similarity: {sim:.4f})" if sim is not None else "" + section = f"[{i}] PMCID: {pmcid}{sim_str}\nTitle: {title}\n" + if abstract: + section += f"Abstract: {abstract}\n" + if full_text: + remaining = max_chars - total - len(section) - 200 + if remaining > 500: + section += f"Full text excerpt: {full_text[:6000][:remaining]}\n" + if total + len(section) > max_chars: + break + parts.append(section) + total += len(section) + return "\n".join(parts) + + +def answer_question(question: str, context: str, llm_client: AzureOpenAI, + model: str, system_prompt: str, temperature: float, + max_completion_tokens: int, max_retries: int) -> str: + prompt = f"""Based on the following retrieved PubMed articles, answer this question: + +Question: {question} + +Retrieved Articles: +{context} + +Provide a comprehensive answer synthesizing the information from these articles.""" + for attempt in range(max_retries): + try: + resp = llm_client.chat.completions.create( + model=model, + messages=[{"role": "system", "content": system_prompt}, + {"role": "user", "content": prompt}], + temperature=temperature, max_completion_tokens=max_completion_tokens, + ) + return resp.choices[0].message.content.strip() + except Exception as e: + log.warning(f"LLM error (attempt {attempt+1}): {e}") + if attempt < max_retries - 1: + time.sleep(min(5.0 * (2 ** attempt), 60)) + else: + return f"Error: {e}" + + +def main(): + pa = argparse.ArgumentParser(description="Answer questions: vector + reference expansion + closest selection") + pa.add_argument("--config", default="config.pubmed.yaml", help="Path to YAML config file") + pa.add_argument("--questions", default=None) + pa.add_argument("--output", default=None) + pa.add_argument("--top", "-k", type=int, default=None, help="Initial vector search docs") + pa.add_argument("--ref-k", type=int, default=None, help="Closest reference docs per hit") + args = pa.parse_args() + + with open(args.config) as f: + cfg = yaml.safe_load(f) + cosmos_cfg = cfg.get("cosmos", {}) + llm_cfg = cfg.get("llm", {}) + embed_cfg = cfg.get("embedding", {}) + answer_cfg = cfg.get("answer", {}) + + # Resolve parameters: CLI overrides > config > defaults + questions_path = args.questions or answer_cfg.get("questions", "questions.json") + output_path = args.output or answer_cfg.get("output", "answers_refs.json") + top_k = args.top if args.top is not None else int(answer_cfg.get("top", 10)) + ref_k = args.ref_k if args.ref_k is not None else int(answer_cfg.get("ref_k", 2)) + + cosmos_uri = cosmos_cfg.get("uri", "") + database_name = cosmos_cfg.get("database_name", "pubmed") + container_name = cosmos_cfg.get("sources", [{}])[0].get("container_name", "articles") + llm_endpoint = llm_cfg.get("endpoint", "") + llm_model = llm_cfg.get("model", "gpt-5.4") + llm_api_version = llm_cfg.get("api_version", "2024-12-01-preview") + llm_token_scope = llm_cfg.get("token_scope", "https://cognitiveservices.azure.com/.default") + llm_temperature = float(llm_cfg.get("temperature", 0.0)) + llm_max_tokens = int(llm_cfg.get("max_completion_tokens", 4096)) + llm_max_retries = int(llm_cfg.get("max_retries", 5)) + embed_endpoint = embed_cfg.get("endpoint", "") + embed_api_key = embed_cfg.get("api_key", "") + embed_api_version = embed_cfg.get("api_version", "2024-12-01-preview") + embed_model = embed_cfg.get("model", "text-embedding-3-small") + embed_dims = int(embed_cfg.get("dimensions", 1536)) + embed_max_input_chars = int(embed_cfg.get("max_input_chars", 8000)) + pmc_ids_csv = answer_cfg.get("pmc_ids_csv", "") + max_context_chars = int(answer_cfg.get("max_context_chars", 230000)) + system_prompt = answer_cfg.get("system_prompt", "You are a biomedical research assistant.").strip() + + with open(questions_path) as f: + raw = json.load(f) + questions = [item["question"] for item in raw] if raw and isinstance(raw[0], dict) else raw + log.info(f"Loaded {len(questions)} questions") + + pmid_map = load_pmid_to_pmcid(pmc_ids_csv) if pmc_ids_csv else {} + log.info(f"Loaded {len(pmid_map)} PMID→PMCID mappings") + + cred = AzureCliCredential() + container = CosmosClient(cosmos_uri, credential=cred) \ + .get_database_client(database_name).get_container_client(container_name) + embed_client = AzureOpenAI( + azure_endpoint=embed_endpoint, api_key=embed_api_key, api_version=embed_api_version) + llm_client = AzureOpenAI( + azure_endpoint=llm_endpoint, + azure_ad_token_provider=get_bearer_token_provider(cred, llm_token_scope), + api_version=llm_api_version) + log.info(f"Strategy: {top_k} vector docs, {ref_k} closest refs each → max {top_k + top_k * ref_k} docs") + + results = [] + for i, question in enumerate(tqdm(questions, desc="Processing")): + log.info(f"[{i+1}/{len(questions)}] {question[:80]}...") + query_vec = get_embedding(question, embed_client, embed_model, embed_dims, embed_max_input_chars) + + # 1. Vector search + hits = vector_search(container, query_vec, top=top_k) + seen_pmcids = {d.get("pmcid") for d in hits} + log.info(f" Retrieved {len(hits)} hits") + + # 2. For each hit, get reference PMCIDs, fetch them, pick ref_k closest to query + ref_docs_all = [] + if ref_k > 0: + for doc in hits: + refs = doc.get("references") or [] + ref_pmcids = [p for p in refs_to_pmcids(refs, pmid_map) if p not in seen_pmcids] + if not ref_pmcids: + continue + fetched = fetch_docs_by_pmcids(container, ref_pmcids) + if not fetched: + continue + selected = closest_select(fetched, query_vec, ref_k) + for s in selected: + if s.get("pmcid") not in seen_pmcids: + seen_pmcids.add(s.get("pmcid")) + ref_docs_all.append(s) + + all_docs = hits + ref_docs_all + log.info(f" Total docs: {len(hits)} hits + {len(ref_docs_all)} ref expansions = {len(all_docs)}") + + # 3. Answer + context = build_context(all_docs, max_chars=max_context_chars) + answer = answer_question(question, context, llm_client, + model=llm_model, system_prompt=system_prompt, + temperature=llm_temperature, + max_completion_tokens=llm_max_tokens, + max_retries=llm_max_retries) + + # 4. Collect + retrieved = [{"pmcid": d.get("pmcid"), "title": d.get("title"), + "journal": d.get("journal_title"), "year": d.get("pub_year"), + "similarity": d.get("similarity"), "source": "vector" if d in hits else "reference"} + for d in all_docs] + results.append({"question": question, "answer": answer, "retrieved_documents": retrieved, + "vector_count": len(hits), "ref_count": len(ref_docs_all)}) + + with open(output_path, "w") as f: + json.dump(results, f, indent=2, ensure_ascii=False) + + log.info(f"Done. Saved {len(results)} answers to {output_path}") + + +if __name__ == "__main__": + main() diff --git a/pubmed/scripts/evaluate_answers_ragas.py b/pubmed/scripts/evaluate_answers_ragas.py new file mode 100644 index 0000000..f7c37e5 --- /dev/null +++ b/pubmed/scripts/evaluate_answers_ragas.py @@ -0,0 +1,219 @@ +"""Compare two answer JSON files using RAGAS metrics. + +Evaluates each answer set independently on multiple criteria, then prints +a side-by-side comparison. Uses Azure OpenAI via litellm + Azure AD token. + +Usage: + python evaluate_answers_ragas.py \ + --base output/review_answers_64.json \ + --other output/review_answers_divdet_64_128_rp5_eta01.json \ + --base-name vector_64 \ + --other-name divdet_64 \ + --output output/eval_ragas_vector_64_vs_divdet_64.json +""" + +import argparse +import asyncio +import json +import logging +import sys +import time + +import litellm +import yaml +from azure.identity import AzureCliCredential +from ragas.llms import llm_factory +from ragas.metrics import NumericMetric +from tqdm import tqdm + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +log = logging.getLogger(__name__) + +# Suppress noisy loggers +for name in ["azure", "httpx", "openai", "litellm", "instructor", "httpcore"]: + logging.getLogger(name).setLevel(logging.WARNING) + +# --------------------------------------------------------------------------- +# Metrics definitions +# --------------------------------------------------------------------------- +METRICS = { + "comprehensiveness": NumericMetric( + name="comprehensiveness", + prompt=( + "Rate how comprehensively the answer covers the question on a scale " + "of 0.0 to 1.0. Consider breadth of coverage, depth of detail, and " + "whether important aspects are missing.\n\n" + "Question: {question}\nAnswer: {answer}" + ), + ), + "evidence_use": NumericMetric( + name="evidence_use", + prompt=( + "Rate how well the answer uses evidence from retrieved sources on a " + "scale of 0.0 to 1.0. Consider citation of specific papers, use of " + "concrete data/findings, and grounding of claims.\n\n" + "Question: {question}\nAnswer: {answer}" + ), + ), + "accuracy": NumericMetric( + name="accuracy", + prompt=( + "Rate the factual accuracy and scientific rigor of the answer on a " + "scale of 0.0 to 1.0. Consider whether claims are well-supported, " + "nuanced, and free of contradictions or errors.\n\n" + "Question: {question}\nAnswer: {answer}" + ), + ), + "clarity": NumericMetric( + name="clarity", + prompt=( + "Rate the clarity and organization of the answer on a scale of 0.0 " + "to 1.0. Consider logical structure, readability, effective use of " + "headings/sections, and overall coherence.\n\n" + "Question: {question}\nAnswer: {answer}" + ), + ), + "relevance": NumericMetric( + name="relevance", + prompt=( + "Rate how relevant the answer is to the specific question asked on " + "a scale of 0.0 to 1.0. Consider whether it directly addresses the " + "question without excessive tangential content.\n\n" + "Question: {question}\nAnswer: {answer}" + ), + ), +} + + +def _make_llm(cred, endpoint, deployment, api_version, token_scope): + """Create a fresh ragas LLM with a new token.""" + token = cred.get_token(token_scope).token + return llm_factory( + f"azure/{deployment}", + provider="litellm", + client=litellm.completion, + api_base=endpoint, + api_version=api_version, + api_key=token, + ) + + +def build_llm(llm_cfg): + """Create a ragas LLM using Azure AD token via litellm.""" + endpoint = llm_cfg.get("endpoint", "") + deployment = llm_cfg.get("model", "gpt-5.4") + api_version = llm_cfg.get("api_version", "2024-12-01-preview") + token_scope = llm_cfg.get("token_scope", "https://cognitiveservices.azure.com/.default") + cred = AzureCliCredential() + llm = _make_llm(cred, endpoint, deployment, api_version, token_scope) + return llm, cred, {"endpoint": endpoint, "deployment": deployment, + "api_version": api_version, "token_scope": token_scope} + + +def load_answers(path: str) -> list[dict]: + with open(path) as f: + return json.load(f) + + +def evaluate_answers(answers: list[dict], llm, cred, llm_params: dict, name: str) -> list[dict]: + """Score each answer on all metrics. Returns list of per-question result dicts.""" + results = [] + for i, item in enumerate(tqdm(answers, desc=f"Scoring [{name}]")): + question = item.get("question", "") + answer = item.get("answer", "") + row = {"question_idx": i, "question": question[:120]} + + # Rebuild LLM with fresh token every 10 questions to avoid expiry + if i % 10 == 0 and i > 0: + llm = _make_llm(cred, **llm_params) + log.info("Rebuilt LLM with fresh token") + + for metric_name, metric in METRICS.items(): + for attempt in range(5): + try: + result = metric.score(llm=llm, question=question, answer=answer) + row[metric_name] = result.value + break + except Exception as e: + err_str = str(e) + if "Authentication" in err_str or "401" in err_str or "403" in err_str: + log.warning(f"[{name}] q{i} {metric_name} auth error, rebuilding LLM...") + llm = _make_llm(cred, **llm_params) + if attempt < 4: + wait = min(5.0 * (2 ** attempt), 60) + log.warning(f"[{name}] q{i} {metric_name} attempt {attempt+1} failed: {err_str[:100]}. Retrying in {wait:.0f}s...") + time.sleep(wait) + else: + log.error(f"[{name}] q{i} {metric_name} failed after 5 attempts: {e}") + row[metric_name] = None + + results.append(row) + if (i + 1) % 10 == 0: + log.info(f"[{name}] Scored {i+1}/{len(answers)} questions") + + log.info(f"[{name}] Done scoring {len(answers)} questions") + return results + + +def main(): + parser = argparse.ArgumentParser(description="Compare two answer files using RAGAS metrics") + parser.add_argument("--config", default="config.pubmed.yaml", help="Path to YAML config file") + parser.add_argument("--base", required=True, help="Base answer JSON file") + parser.add_argument("--other", required=True, help="Other answer JSON file") + parser.add_argument("--base-name", default="base", help="Name for base method") + parser.add_argument("--other-name", default="other", help="Name for other method") + parser.add_argument("--output", "-o", default=None, help="Output JSON file") + args = parser.parse_args() + + with open(args.config) as f: + cfg = yaml.safe_load(f) + llm_cfg = cfg.get("llm", {}) + + base_answers = load_answers(args.base) + other_answers = load_answers(args.other) + log.info(f"Loaded {len(base_answers)} from {args.base}, {len(other_answers)} from {args.other}") + + assert len(base_answers) == len(other_answers), "Answer files must have the same number of questions" + + llm, cred, llm_params = build_llm(llm_cfg) + + log.info(f"Evaluating '{args.base_name}'...") + base_results = evaluate_answers(base_answers, llm, cred, llm_params, args.base_name) + + log.info(f"Evaluating '{args.other_name}'...") + other_results = evaluate_answers(other_answers, llm, cred, llm_params, args.other_name) + + # Combine results + combined = [] + for b, o in zip(base_results, other_results): + row = {"question_idx": b["question_idx"], "question": b["question"]} + for metric_name in METRICS: + row[f"{args.base_name}_{metric_name}"] = b.get(metric_name) + row[f"{args.other_name}_{metric_name}"] = o.get(metric_name) + combined.append(row) + + # Save + output_path = args.output or f"eval_ragas_{args.base_name}_vs_{args.other_name}.json" + with open(output_path, "w") as f: + json.dump(combined, f, indent=2, ensure_ascii=False) + log.info(f"Saved results to {output_path}") + + # Print summary + print(f"\n{'='*70}") + print(f"RAGAS Comparison: {args.base_name} vs {args.other_name}") + print(f"{'='*70}") + print(f"{'Metric':<25} {args.base_name:>15} {args.other_name:>15} {'delta':>10}") + print(f"{'-'*70}") + for metric_name in METRICS: + base_vals = [r.get(f"{args.base_name}_{metric_name}") for r in combined if r.get(f"{args.base_name}_{metric_name}") is not None] + other_vals = [r.get(f"{args.other_name}_{metric_name}") for r in combined if r.get(f"{args.other_name}_{metric_name}") is not None] + base_mean = sum(base_vals) / len(base_vals) if base_vals else 0 + other_mean = sum(other_vals) / len(other_vals) if other_vals else 0 + delta = other_mean - base_mean + sign = "+" if delta >= 0 else "" + print(f"{metric_name:<25} {base_mean:>15.4f} {other_mean:>15.4f} {sign}{delta:>9.4f}") + print(f"{'='*70}") + + +if __name__ == "__main__": + main() diff --git a/pubmed/scripts/readme.md b/pubmed/scripts/readme.md new file mode 100644 index 0000000..d6f1523 --- /dev/null +++ b/pubmed/scripts/readme.md @@ -0,0 +1,38 @@ +Run script without using graph augmentation (references): + +``` +python answer_questions_ref.py --config config.pubmed.yaml --questions ../data/complex_questions.json --output output/answers/answers_top15.json --top 15 --ref-k 0 +``` + +This script will produce an answer file where each answer is generated from 15 documents. + +Run script with graph augmentation (references): + +``` +python answer_questions_ref.py --config config.pubmed.yaml --questions ../data/complex_questions.json --output output/answers_ref/answers_top5_refk4.json --top 5 --ref-k 4 +``` + +The last script will produce an answer file where each answer is generated from 16.78 documents on average (the average over 100 questions). + +Now we can compare the 2 produced answer files using ragas: + +``` +python evaluate_answers_ragas.py --base output/answers/answers_top15.json --other output/answers_ref/answers_top5_refk4.json --base-name vector15 --other-name refs_5_4 --output output/eval_ragas.json +``` + +This should produce comparison similar to this: + + +
+=====================================================================
+RAGAS Comparison: vector15 vs refs_5_4
+======================================================================
+Metric                           vector15        refs_5_4      delta
+----------------------------------------------------------------------
+comprehensiveness                  0.9278          0.9240   -0.0038
+evidence_use                       0.8940          0.9060 +   0.0120
+accuracy                           0.8908          0.8899   -0.0009
+clarity                            0.9329          0.9349 +   0.0020
+relevance                          0.9429          0.9414   -0.0015
+======================================================================
+
\ No newline at end of file diff --git a/utils/cosmos_retriever.py b/utils/cosmos_retriever.py index d67ca99..94190a6 100644 --- a/utils/cosmos_retriever.py +++ b/utils/cosmos_retriever.py @@ -91,8 +91,9 @@ def _get_source_config(config: dict[str, Any]) -> list[dict[str, Any]]: RETRIEVAL_SOURCES = _get_source_config(CONFIG) -# Comprehensive BM25 stopwords list -STOPWORDS = {"0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "a", "a's", "able", "about", "above", "according", "accordingly", "across", "actually", "after", "afterwards", "again", "against", "ain't", "all", "allow", "allows", "almost", "alone", "along", "already", "also", "although", "always", "am", "among", "amongst", "an", "and", "another", "any", "anybody", "anyhow", "anyone", "anything", "anyway", "anyways", "anywhere", "apart", "appear", "appreciate", "appropriate", "are", "aren't", "around", "as", "aside", "ask", "asking", "associated", "at", "available", "away", "awfully", "b", "be", "became", "because", "become", "becomes", "becoming", "been", "before", "beforehand", "behind", "being", "believe", "below", "beside", "besides", "best", "better", "between", "beyond", "both", "brief", "but", "by", "c", "c'mon", "c's", "came", "can", "can't", "cannot", "cant", "cause", "causes", "certain", "certainly", "changes", "clearly", "co", "com", "come", "comes", "concerning", "consequently", "consider", "considering", "contain", "containing", "contains", "corresponding", "could", "couldn't", "course", "currently", "d", "definitely", "described", "despite", "did", "didn't", "different", "do", "does", "doesn't", "doing", "don", "don't", "done", "down", "downwards", "during", "e", "each", "edu", "eg", "eight", "either", "else", "elsewhere", "enough", "entirely", "especially", "et", "etc", "even", "ever", "every", "everybody", "everyone", "everything", "everywhere", "ex", "exactly", "example", "except", "f", "far", "few", "fifth", "first", "five", "followed", "following", "follows", "for", "former", "formerly", "forth", "four", "from", "further", "furthermore", "g", "get", "gets", "getting", "given", "gives", "go", "goes", "going", "gone", "got", "gotten", "greetings", "h", "had", "hadn't", "happens", "hardly", "has", "hasn't", "have", "haven't", "having", "he", "he's", "hello", "help", "hence", "her", "here", "here's", "hereafter", "hereby", "herein", "hereupon", "hers", "herself", "hi", "him", "himself", "his", "hither", "hopefully", "how", "howbeit", "however", "i", "i'd", "i'll", "i'm", "i've", "ie", "if", "ignored", "immediate", "in", "inasmuch", "inc", "indeed", "indicate", "indicated", "indicates", "inner", "insofar", "instead", "into", "inward", "is", "isn't", "it", "it'd", "it'll", "it's", "its", "itself", "j", "just", "k", "keep", "keeps", "kept", "know", "known", "knows", "l", "last", "lately", "later", "latter", "latterly", "least", "less", "lest", "let", "let's", "like", "liked", "likely", "little", "ll", "look", "looking", "looks", "ltd", "m", "mainly", "make", "many", "may", "maybe", "me", "mean", "meanwhile", "merely", "might", "more", "moreover", "most", "mostly", "mr", "mrs", "ms", "much", "must", "my", "myself", "n", "name", "namely", "nd", "near", "nearly", "necessary", "need", "needs", "neither", "never", "nevertheless", "new", "next", "nine", "no", "nobody", "non", "none", "noone", "nor", "normally", "not", "nothing", "novel", "now", "nowhere", "o", "obviously", "of", "off", "often", "oh", "ok", "okay", "old", "on", "once", "one", "ones", "only", "onto", "or", "other", "others", "otherwise", "ought", "our", "ours", "ourselves", "out", "outside", "over", "overall", "own", "p", "particular", "particularly", "per", "perhaps", "placed", "please", "plus", "possible", "presumably", "probably", "provides", "q", "que", "quite", "qv", "r", "rather", "rd", "re", "really", "reasonably", "regarding", "regardless", "regards", "relatively", "respectively", "right", "s", "said", "same", "saw", "say", "saying", "says", "second", "secondly", "see", "seeing", "seem", "seemed", "seeming", "seems", "seen", "self", "selves", "sensible", "sent", "serious", "seriously", "seven", "several", "shall", "she", "should", "shouldn't", "since", "six", "so", "some", "somebody", "somehow", "someone", "something", "sometime", "sometimes", "somewhat", "somewhere", "soon", "sorry", "specified", "specify", "specifying", "still", "sub", "such", "sup", "sure", "t", "t's", "take", "taken", "tell", "tends", "th", "than", "thank", "thanks", "thanx", "that", "that's", "thats", "the", "their", "theirs", "them", "themselves", "then", "thence", "there", "there's", "thereafter", "thereby", "therefore", "therein", "theres", "thereupon", "these", "they", "they'd", "they'll", "they're", "they've", "think", "third", "this", "thorough", "thoroughly", "those", "though", "three", "through", "throughout", "thru", "thus", "to", "together", "too", "took", "toward", "towards", "tried", "tries", "truly", "try", "trying", "twice", "two", "u", "un", "under", "unfortunately", "unless", "unlikely", "until", "unto", "up", "upon", "us", "use", "used", "useful", "uses", "using", "usually", "v", "value", "various", "ve", "very", "via", "viz", "vs", "w", "want", "wants", "was", "wasn't", "way", "we", "we'd", "we'll", "we're", "we've", "welcome", "well", "went", "were", "weren't", "what", "what's", "whatever", "when", "whence", "whenever", "where", "where's", "whereafter", "whereas", "whereby", "wherein", "whereupon", "wherever", "whether", "which", "while", "whither", "who", "who's", "whoever", "whole", "whom", "whose", "why", "will", "willing", "wish", "with", "within", "without", "won't", "wonder", "would", "wouldn't", "x", "y", "yes", "yet", "you", "you'd", "you'll", "you're", "you've", "your", "yours", "yourself", "yourselves", "z", "zero"} +# Re-export from utils.fulltext for backward compatibility +from utils.fulltext import fulltext_search # noqa: E402 +from utils.ranker import rerank_documents # noqa: E402 class CombinedRetriever: @@ -248,85 +249,13 @@ async def initialize(self): except Exception as e: _log_line(f"Ranker account registration failed: {e}", kind="warn") - async def _fulltext_search_single_field(self, container, field: str, query: str, top_k: int) -> list[dict]: - """Run a fulltext search on a single field, returning ranked results.""" - if top_k <= 0: - return [] - terms = [t for t in re.findall(r"\w+", query) if t.lower() not in STOPWORDS and len(t) > 1] - if not terms: - return [] - - chunks = [terms[i:i + 5] for i in range(0, len(terms), 5)] - field_expr = f"c.{field}" - score_exprs = [] - for term_chunk in chunks: - args = ", ".join(f'"{term}"' for term in term_chunk) - score_exprs.append(f"FullTextScore({field_expr}, {args})") - if not score_exprs: - return [] - - if len(score_exprs) == 1: - order = f"ORDER BY RANK {score_exprs[0]}" - else: - order = f"ORDER BY RANK RRF({', '.join(score_exprs)})" - - sql = f"SELECT TOP {top_k} * FROM c {order}" - try: - if _timing_enabled(): - _log_line( - f" fulltext SQL ({container.id}/{field}): {sql} [text={query!r}]", - kind="query", - use_lock=True, - ) - - t = _ck(f"fulltext query (top {top_k}, {container.id}/{field}) – start") - query_iterator = container.query_items(query=sql, parameters=[]) - items = [] - async for item in query_iterator: - items.append(item) - - _ck(f"fulltext query – done ({len(items)} results, {container.id}/{field})", t) - return items - except Exception as e: - _log_line(f"Fulltext error ({container.id}/{field}): {e}", kind="error") - return [] - async def _fulltext_search(self, container, fields: list[str], query: str, top_k: int) -> list[dict]: if top_k <= 0 or not fields: return [] - terms = [t for t in re.findall(r"\w+", query) if t.lower() not in STOPWORDS and len(t) > 1] - if not terms: - return [] - - # Single field: run directly (no need for client-side RRF) - if len(fields) == 1: - return await self._fulltext_search_single_field(container, fields[0], query, top_k) - - # Multiple fields: run parallel per-field queries and merge via client-side RRF - t = _ck(f"fulltext parallel-RRF (top {top_k}, {container.id}, {len(fields)} fields) – start") - per_field_tasks = [ - asyncio.create_task(self._fulltext_search_single_field(container, field, query, top_k)) - for field in fields - ] - per_field_results = await asyncio.gather(*per_field_tasks) - - # Client-side RRF merge: score each doc by 1/(k+rank) across fields, pick top_k - rrf_k = 60 # standard RRF constant - doc_scores: dict[str, float] = {} - doc_map: dict[str, dict] = {} - for field_items in per_field_results: - for rank, item in enumerate(field_items): - doc_id = item.get("id", "") - if not doc_id: - continue - doc_scores[doc_id] = doc_scores.get(doc_id, 0.0) + 1.0 / (rrf_k + rank + 1) - if doc_id not in doc_map: - doc_map[doc_id] = item - - sorted_ids = sorted(doc_scores, key=lambda did: doc_scores[did], reverse=True)[:top_k] - merged = [doc_map[did] for did in sorted_ids] - _ck(f"fulltext parallel-RRF – done ({len(merged)} merged from {sum(len(r) for r in per_field_results)} total, {container.id})", t) - return merged + t = _ck(f"fulltext (top {top_k}, {container.id}, {len(fields)} fields) – start") + items = await fulltext_search(container, fields, query, top_k) + _ck(f"fulltext – done ({len(items)} results, {container.id})", t) + return items async def _vector_search( self, @@ -540,47 +469,18 @@ async def retrieve(self, query: str, k_divisor: int = 1) -> list[RetrievedChunk] if self._ranker_http_client is None: self._ranker_http_client = httpx.AsyncClient(timeout=120) documents = [c.text for c in chunks] - body = { - "query": query, - "documents": documents, - "return_documents": False, - "top_k": effective_k_ranker, - "batch_size": self._ranker_batch_size, - } - headers = { - "Authorization": f"Bearer {self._ranker_access_token}", - "Content-Type": "application/json", - } url_suffix = str(CONFIG.get("ranker", {}).get("url_suffix", "")).strip() url = f"https://{self._ranker_account}.{self._ranker_region}.{url_suffix}" max_retries = int(CONFIG.get("ranker", {}).get("max_retries", 5)) - ranker_succeeded = False - for attempt in range(max_retries): - try: - response = await self._ranker_http_client.post(url, headers=headers, json=body) - if response.status_code in (502, 503, 429) and attempt < max_retries - 1: - wait = 2 ** attempt - _log_line(f"Semantic ranker returned {response.status_code}, retrying in {wait}s (attempt {attempt + 1}/{max_retries})", kind="warn") - await asyncio.sleep(wait) - continue - response.raise_for_status() - result = response.json() - scores = result.get("Scores", []) - # Select top k_ranker by reranker score, preserving original chunk objects - ranked_indices = [s["index"] for s in scores[:effective_k_ranker]] - chunks = [chunks[i] for i in ranked_indices] - _ck(f" retrieve: semantic ranker – done (selected {len(chunks)} of {effective_k_ranker} requested)", t) - ranker_succeeded = True - break - except Exception as e: - if attempt < max_retries - 1 and ("503" in str(e) or "502" in str(e) or "429" in str(e)): - wait = 2 ** attempt - _log_line(f"Semantic ranker error (attempt {attempt + 1}/{max_retries}): {e}, retrying in {wait}s", kind="warn") - await asyncio.sleep(wait) - continue - _log_line(f"Semantic ranker error: {e}", kind="error") - break - if not ranker_succeeded: + ranked_indices = await rerank_documents( + self._ranker_http_client, url, self._ranker_access_token, + query, documents, effective_k_ranker, self._ranker_batch_size, max_retries, + ) + if ranked_indices is not None: + chunks = [chunks[i] for i in ranked_indices] + _ck(f" retrieve: semantic ranker – done (selected {len(chunks)} of {effective_k_ranker} requested)", t) + else: + _log_line("Semantic ranker failed, keeping previous chunks", kind="error") _ck(" retrieve: semantic ranker – failed, keeping diversity-selected chunks", t) self._retrieve_cache.set(cache_key, copy.deepcopy(chunks)) diff --git a/utils/fulltext.py b/utils/fulltext.py new file mode 100644 index 0000000..1efee84 --- /dev/null +++ b/utils/fulltext.py @@ -0,0 +1,69 @@ +"""Shared fulltext search helpers and stopwords. + +This module has no heavy dependencies (no CONFIG, no agentic_retriever) so it +can be imported safely from both cosmos_retriever.py and dynamic_retriever.py. +""" + +import asyncio +import re + +# Comprehensive BM25 stopwords list +STOPWORDS = {"0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "a", "a's", "able", "about", "above", "according", "accordingly", "across", "actually", "after", "afterwards", "again", "against", "ain't", "all", "allow", "allows", "almost", "alone", "along", "already", "also", "although", "always", "am", "among", "amongst", "an", "and", "another", "any", "anybody", "anyhow", "anyone", "anything", "anyway", "anyways", "anywhere", "apart", "appear", "appreciate", "appropriate", "are", "aren't", "around", "as", "aside", "ask", "asking", "associated", "at", "available", "away", "awfully", "b", "be", "became", "because", "become", "becomes", "becoming", "been", "before", "beforehand", "behind", "being", "believe", "below", "beside", "besides", "best", "better", "between", "beyond", "both", "brief", "but", "by", "c", "c'mon", "c's", "came", "can", "can't", "cannot", "cant", "cause", "causes", "certain", "certainly", "changes", "clearly", "co", "com", "come", "comes", "concerning", "consequently", "consider", "considering", "contain", "containing", "contains", "corresponding", "could", "couldn't", "course", "currently", "d", "definitely", "described", "despite", "did", "didn't", "different", "do", "does", "doesn't", "doing", "don", "don't", "done", "down", "downwards", "during", "e", "each", "edu", "eg", "eight", "either", "else", "elsewhere", "enough", "entirely", "especially", "et", "etc", "even", "ever", "every", "everybody", "everyone", "everything", "everywhere", "ex", "exactly", "example", "except", "f", "far", "few", "fifth", "first", "five", "followed", "following", "follows", "for", "former", "formerly", "forth", "four", "from", "further", "furthermore", "g", "get", "gets", "getting", "given", "gives", "go", "goes", "going", "gone", "got", "gotten", "greetings", "h", "had", "hadn't", "happens", "hardly", "has", "hasn't", "have", "haven't", "having", "he", "he's", "hello", "help", "hence", "her", "here", "here's", "hereafter", "hereby", "herein", "hereupon", "hers", "herself", "hi", "him", "himself", "his", "hither", "hopefully", "how", "howbeit", "however", "i", "i'd", "i'll", "i'm", "i've", "ie", "if", "ignored", "immediate", "in", "inasmuch", "inc", "indeed", "indicate", "indicated", "indicates", "inner", "insofar", "instead", "into", "inward", "is", "isn't", "it", "it'd", "it'll", "it's", "its", "itself", "j", "just", "k", "keep", "keeps", "kept", "know", "known", "knows", "l", "last", "lately", "later", "latter", "latterly", "least", "less", "lest", "let", "let's", "like", "liked", "likely", "little", "ll", "look", "looking", "looks", "ltd", "m", "mainly", "make", "many", "may", "maybe", "me", "mean", "meanwhile", "merely", "might", "more", "moreover", "most", "mostly", "mr", "mrs", "ms", "much", "must", "my", "myself", "n", "name", "namely", "nd", "near", "nearly", "necessary", "need", "needs", "neither", "never", "nevertheless", "new", "next", "nine", "no", "nobody", "non", "none", "noone", "nor", "normally", "not", "nothing", "novel", "now", "nowhere", "o", "obviously", "of", "off", "often", "oh", "ok", "okay", "old", "on", "once", "one", "ones", "only", "onto", "or", "other", "others", "otherwise", "ought", "our", "ours", "ourselves", "out", "outside", "over", "overall", "own", "p", "particular", "particularly", "per", "perhaps", "placed", "please", "plus", "possible", "presumably", "probably", "provides", "q", "que", "quite", "qv", "r", "rather", "rd", "re", "really", "reasonably", "regarding", "regardless", "regards", "relatively", "respectively", "right", "s", "said", "same", "saw", "say", "saying", "says", "second", "secondly", "see", "seeing", "seem", "seemed", "seeming", "seems", "seen", "self", "selves", "sensible", "sent", "serious", "seriously", "seven", "several", "shall", "she", "should", "shouldn't", "since", "six", "so", "some", "somebody", "somehow", "someone", "something", "sometime", "sometimes", "somewhat", "somewhere", "soon", "sorry", "specified", "specify", "specifying", "still", "sub", "such", "sup", "sure", "t", "t's", "take", "taken", "tell", "tends", "th", "than", "thank", "thanks", "thanx", "that", "that's", "thats", "the", "their", "theirs", "them", "themselves", "then", "thence", "there", "there's", "thereafter", "thereby", "therefore", "therein", "theres", "thereupon", "these", "they", "they'd", "they'll", "they're", "they've", "think", "third", "this", "thorough", "thoroughly", "those", "though", "three", "through", "throughout", "thru", "thus", "to", "together", "too", "took", "toward", "towards", "tried", "tries", "truly", "try", "trying", "twice", "two", "u", "un", "under", "unfortunately", "unless", "unlikely", "until", "unto", "up", "upon", "us", "use", "used", "useful", "uses", "using", "usually", "v", "value", "various", "ve", "very", "via", "viz", "vs", "w", "want", "wants", "was", "wasn't", "way", "we", "we'd", "we'll", "we're", "we've", "welcome", "well", "went", "were", "weren't", "what", "what's", "whatever", "when", "whence", "whenever", "where", "where's", "whereafter", "whereas", "whereby", "wherein", "whereupon", "wherever", "whether", "which", "while", "whither", "who", "who's", "whoever", "whole", "whom", "whose", "why", "will", "willing", "wish", "with", "within", "without", "won't", "wonder", "would", "wouldn't", "x", "y", "yes", "yet", "you", "you'd", "you'll", "you're", "you've", "your", "yours", "yourself", "yourselves", "z", "zero"} + + +async def fulltext_search_single_field(container, field: str, query: str, top_k: int) -> list[dict]: + """Run a fulltext search on a single field, returning ranked results.""" + if top_k <= 0: + return [] + terms = [t for t in re.findall(r"\w+", query) if t.lower() not in STOPWORDS and len(t) > 1] + if not terms: + return [] + chunks = [terms[i:i + 5] for i in range(0, len(terms), 5)] + field_expr = f"c.{field}" + score_exprs = [] + for term_chunk in chunks: + args = ", ".join(f'"{term}"' for term in term_chunk) + score_exprs.append(f"FullTextScore({field_expr}, {args})") + if not score_exprs: + return [] + if len(score_exprs) == 1: + order = f"ORDER BY RANK {score_exprs[0]}" + else: + order = f"ORDER BY RANK RRF({', '.join(score_exprs)})" + sql = f"SELECT TOP {top_k} * FROM c {order}" + try: + items = [] + async for item in container.query_items(query=sql, parameters=[]): + items.append(item) + return items + except asyncio.CancelledError: + raise + except Exception: + return [] + + +async def fulltext_search(container, fields: list[str], query: str, top_k: int) -> list[dict]: + """Multi-field fulltext search with client-side RRF merge.""" + if top_k <= 0 or not fields: + return [] + terms = [t for t in re.findall(r"\w+", query) if t.lower() not in STOPWORDS and len(t) > 1] + if not terms: + return [] + if len(fields) == 1: + return await fulltext_search_single_field(container, fields[0], query, top_k) + per_field_results = await asyncio.gather( + *(fulltext_search_single_field(container, f, query, top_k) for f in fields) + ) + rrf_k = 60 + doc_scores: dict[str, float] = {} + doc_map: dict[str, dict] = {} + for field_items in per_field_results: + for rank, item in enumerate(field_items): + doc_id = item.get("id", "") + if not doc_id: + continue + doc_scores[doc_id] = doc_scores.get(doc_id, 0.0) + 1.0 / (rrf_k + rank + 1) + if doc_id not in doc_map: + doc_map[doc_id] = item + sorted_ids = sorted(doc_scores, key=lambda did: doc_scores[did], reverse=True)[:top_k] + return [doc_map[did] for did in sorted_ids] diff --git a/utils/ranker.py b/utils/ranker.py new file mode 100644 index 0000000..40079aa --- /dev/null +++ b/utils/ranker.py @@ -0,0 +1,57 @@ +"""Shared semantic ranker helper. + +This module has no heavy dependencies (no CONFIG, no agentic_retriever) so it +can be imported safely from both cosmos_retriever.py and dynamic_retriever.py. +""" + +import asyncio + +import httpx + + +async def rerank_documents( + http_client: httpx.AsyncClient, + url: str, + access_token: str, + query: str, + documents: list[str], + top_k: int, + batch_size: int = 32, + max_retries: int = 5, +) -> list[int] | None: + """Call the semantic ranker and return ranked indices, or None on failure. + + Returns a list of indices into *documents* ordered by ranker score + (best first), or ``None`` if all attempts failed. + """ + if not documents or top_k <= 0: + return list(range(min(len(documents), top_k))) + + body = { + "query": query, + "documents": documents, + "return_documents": False, + "top_k": top_k, + "batch_size": batch_size, + } + headers = { + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json", + } + + for attempt in range(max_retries): + try: + resp = await http_client.post(url, headers=headers, json=body) + if resp.status_code in (429, 502, 503) and attempt + 1 < max_retries: + await asyncio.sleep(2 ** attempt) + continue + resp.raise_for_status() + scores = resp.json().get("Scores", []) + return [s["index"] for s in scores[:top_k] if s["index"] < len(documents)] + except Exception: + if attempt + 1 < max_retries: + await asyncio.sleep(2 ** attempt) + continue + return None + + return None From 5b4705f00250322beb1b5d81e96a8b5056bdac56 Mon Sep 17 00:00:00 2001 From: backurs Date: Wed, 22 Apr 2026 19:11:24 -0700 Subject: [PATCH 2/2] graph augmentation --- pubmed/scripts/answer_questions_ref.py | 12 ++++++------ pubmed/scripts/config.pubmed.yaml | 20 ++++++++++++++++---- pubmed/scripts/evaluate_answers_ragas.py | 4 ++-- pubmed/scripts/readme.md | 1 - 4 files changed, 24 insertions(+), 13 deletions(-) diff --git a/pubmed/scripts/answer_questions_ref.py b/pubmed/scripts/answer_questions_ref.py index 492e78c..898eaae 100644 --- a/pubmed/scripts/answer_questions_ref.py +++ b/pubmed/scripts/answer_questions_ref.py @@ -181,18 +181,18 @@ def main(): cosmos_uri = cosmos_cfg.get("uri", "") database_name = cosmos_cfg.get("database_name", "pubmed") container_name = cosmos_cfg.get("sources", [{}])[0].get("container_name", "articles") - llm_endpoint = llm_cfg.get("endpoint", "") - llm_model = llm_cfg.get("model", "gpt-5.4") + llm_endpoint = llm_cfg.get("llm_endpoint", "") + llm_model = llm_cfg.get("llm_model", "gpt-5.4") llm_api_version = llm_cfg.get("api_version", "2024-12-01-preview") llm_token_scope = llm_cfg.get("token_scope", "https://cognitiveservices.azure.com/.default") llm_temperature = float(llm_cfg.get("temperature", 0.0)) llm_max_tokens = int(llm_cfg.get("max_completion_tokens", 4096)) llm_max_retries = int(llm_cfg.get("max_retries", 5)) - embed_endpoint = embed_cfg.get("endpoint", "") - embed_api_key = embed_cfg.get("api_key", "") + embed_endpoint = embed_cfg.get("embed_endpoint", "") + embed_api_key = embed_cfg.get("embed_api_key", "") embed_api_version = embed_cfg.get("api_version", "2024-12-01-preview") - embed_model = embed_cfg.get("model", "text-embedding-3-small") - embed_dims = int(embed_cfg.get("dimensions", 1536)) + embed_model = embed_cfg.get("embed_model", "text-embedding-3-small") + embed_dims = int(embed_cfg.get("embed_dimensions", 1536)) embed_max_input_chars = int(embed_cfg.get("max_input_chars", 8000)) pmc_ids_csv = answer_cfg.get("pmc_ids_csv", "") max_context_chars = int(answer_cfg.get("max_context_chars", 230000)) diff --git a/pubmed/scripts/config.pubmed.yaml b/pubmed/scripts/config.pubmed.yaml index b8f8311..2f20719 100644 --- a/pubmed/scripts/config.pubmed.yaml +++ b/pubmed/scripts/config.pubmed.yaml @@ -54,11 +54,11 @@ paths: # --------------------------------------------------------------------------- llm: llm_endpoint: "" - api_version: "2024-05-01-preview" - llm_model: "gpt-5.4-mini" + api_version: "2024-12-01-preview" + llm_model: "" temperature: 0.0 max_completion_tokens: 4096 - max_retries: 2 + max_retries: 5 premium_max_concurrency: 4 prompt_cache_size: 2048 use_rbac_auth: true @@ -79,10 +79,11 @@ local_llm: # --------------------------------------------------------------------------- embedding: embed_endpoint: "" - api_version: "2024-05-01-preview" + api_version: "2024-12-01-preview" embed_model: "text-embedding-3-small" embed_dimensions: 1536 embed_cache_size: 4096 + max_input_chars: 8000 use_rbac_auth: true embed_api_key: "" token_scope: "https://cognitiveservices.azure.com/.default" @@ -197,3 +198,14 @@ cosmos: { "path": "/full_text", "language": "en-US" } ] } + +# --------------------------------------------------------------------------- +# Answer generation settings +# --------------------------------------------------------------------------- +answer: + pmc_ids_csv: "data/PMC-ids.csv" + max_context_chars: 230000 + system_prompt: | + You are a biomedical research assistant. You are given a question and a set of retrieved PubMed articles (title, abstract, and potentially full text excerpts). + + Synthesize the information from these articles to provide a comprehensive, evidence-based answer to the question. Cite the papers by their PMCID when referencing specific findings. If the retrieved articles do not contain sufficient information to fully answer the question, state what is covered and what gaps remain. \ No newline at end of file diff --git a/pubmed/scripts/evaluate_answers_ragas.py b/pubmed/scripts/evaluate_answers_ragas.py index f7c37e5..c380e0e 100644 --- a/pubmed/scripts/evaluate_answers_ragas.py +++ b/pubmed/scripts/evaluate_answers_ragas.py @@ -100,8 +100,8 @@ def _make_llm(cred, endpoint, deployment, api_version, token_scope): def build_llm(llm_cfg): """Create a ragas LLM using Azure AD token via litellm.""" - endpoint = llm_cfg.get("endpoint", "") - deployment = llm_cfg.get("model", "gpt-5.4") + endpoint = llm_cfg.get("llm_endpoint", "") + deployment = llm_cfg.get("llm_model", "gpt-5.4") api_version = llm_cfg.get("api_version", "2024-12-01-preview") token_scope = llm_cfg.get("token_scope", "https://cognitiveservices.azure.com/.default") cred = AzureCliCredential() diff --git a/pubmed/scripts/readme.md b/pubmed/scripts/readme.md index d6f1523..2dfdf9f 100644 --- a/pubmed/scripts/readme.md +++ b/pubmed/scripts/readme.md @@ -22,7 +22,6 @@ python evaluate_answers_ragas.py --base output/answers/answers_top15.json --othe This should produce comparison similar to this: -
 =====================================================================
 RAGAS Comparison: vector15 vs refs_5_4