From 989259d30347cacd0c002d05402fca1ec2933b61 Mon Sep 17 00:00:00 2001 From: backurs Date: Tue, 21 Apr 2026 10:02:56 -0700 Subject: [PATCH 1/6] dynamic retriever with native function calling --- config_dynamic.yaml.example | 163 ++++++++++++++++++++++++ dynamic_retriever.py | 240 ++++++++++++++++++++++++++++++++++++ 2 files changed, 403 insertions(+) create mode 100644 config_dynamic.yaml.example create mode 100644 dynamic_retriever.py diff --git a/config_dynamic.yaml.example b/config_dynamic.yaml.example new file mode 100644 index 0000000..2dfec9e --- /dev/null +++ b/config_dynamic.yaml.example @@ -0,0 +1,163 @@ +# Paths (relative to working directory) +paths: + questions_path: "" + output_root: "" + +prune_k: 20 +rerank_multiplier: 8 +context_limit: 270000 + +# 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: "" + +# 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 + 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..e0f6736 --- /dev/null +++ b/dynamic_retriever.py @@ -0,0 +1,240 @@ +"""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 + +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"} + +_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"]) + return [float(x) for x in r.data[0].embedding[:embed_cfg.get("embed_dimensions", 1536)]] + +async def vec_search(container, emb, top_k, ef): + 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 ft_field(container, field, query, top_k): + terms = [t for t in re.findall(r"\w+", query) if t.lower() not in STOPWORDS and len(t) > 1] + if not terms or top_k <= 0: return [] + chunks = [terms[i:i+5] for i in range(0, len(terms), 5)] + exprs = [f'FullTextScore(c.{field}, {", ".join(chr(34)+t.replace(chr(34),"")+chr(34) for t in ch)})' for ch in chunks] + order = f"ORDER BY RANK {exprs[0]}" if len(exprs)==1 else f"ORDER BY RANK RRF({', '.join(exprs)})" + try: return [item async for item in container.query_items(query=f"SELECT TOP {top_k} * FROM c {order}", parameters=[])] + except: return [] + +async def ft_search(container, fields, query, top_k): + if not fields or top_k <= 0: return [] + if len(fields) == 1: return await ft_field(container, fields[0], query, top_k) + pf = await asyncio.gather(*(ft_field(container, f, query, top_k) for f in fields)) + sc, dm = {}, {} + for items in pf: + for r, it in enumerate(items): + d = it.get("id","") + if d: sc[d] = sc.get(d,0) + 1/(60+r+1); dm.setdefault(d, it) + return [dm[d] for d in sorted(sc, key=sc.get, reverse=True)[:top_k]] + +async def rerank(query, docs, top_k): + if not USE_RANKER or not docs: return docs[:top_k] + body = {"query": query, "documents": docs, "return_documents": False, "top_k": top_k, "batch_size": _r_bs} + for att in range(_r_mr): + resp = await _r_http.post(_r_url, headers=_r_hdr, json=body) + if resp.status_code in (429,502,503) and att+1 < _r_mr: await asyncio.sleep(2**att); continue + resp.raise_for_status() + return [docs[s["index"]] for s in resp.json().get("Scores",[])[:top_k] if s["index"] < len(docs)] + return docs + +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(ft_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()) + ranked = await rerank(query, [fmt(d) for d in all_d], total_k) + return json.dumps([{"docid": re.search(r'^id: (.+)$', d, re.MULTILINE).group(1) if re.search(r'^id: (.+)$', d, re.MULTILINE) else "", "snippet": d[:2000]} for d in ranked]) + +async def do_get_doc(docid, containers): + for c in containers.values(): + 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: 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 + for c in containers.values(): + 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'); break + except: 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=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: {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 + async def _exec(t): + a = json.loads(t.function.arguments) + 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: pass + 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: pass + 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}) + print(f" [prune] Kept {len(json.loads(t.function.arguments)['docids'])} docs, context reset") + pruned = True + break + if pruned: + continue + for t, out, _ in results: + msgs.append({"role": "tool", "tool_call_id": t.id, "content": out}) + a = json.loads(t.function.arguments) + print(f" [{t.function.name}] {list(a.values())[0][:80] if isinstance(list(a.values())[0], str) else '...'}") + 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()) + cosmos = CosmosClient(cosmos_cfg["uri"], credential=AsyncAzureCliCredential()) + db = cosmos.get_database_client(cosmos_cfg["database_name"]) + containers = {s["id"]: db.get_container_client(s["container_name"]) for s in sources} + + 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}") + await cosmos.close() + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--config", default="config_dynamic.yaml") + 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_client = AsyncAzureOpenAI(api_version=embed_cfg["api_version"], azure_endpoint=embed_cfg["embed_endpoint"], 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"] + _r_tok = Path(rcfg["access_token_path"]).read_text().strip() if rcfg["read_token_from_path"] else __import__("azure.identity",fromlist=["AzureCliCredential"]).AzureCliCredential(tenant_id=rcfg["tenant_id"]).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()) From 21a87cf679e0b0afc3ae164b46a7a3d2b945de27 Mon Sep 17 00:00:00 2001 From: backurs Date: Tue, 21 Apr 2026 10:59:23 -0700 Subject: [PATCH 2/6] fix example yaml file --- config_dynamic.yaml.example | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/config_dynamic.yaml.example b/config_dynamic.yaml.example index 2dfec9e..38bb959 100644 --- a/config_dynamic.yaml.example +++ b/config_dynamic.yaml.example @@ -4,8 +4,6 @@ paths: output_root: "" prune_k: 20 -rerank_multiplier: 8 -context_limit: 270000 # Azure OpenAI / LLM settings llm: @@ -20,6 +18,7 @@ llm: use_rbac_auth: true token_scope: "" llm_api_key: "" + context_limit: 270000 # Embedding settings embedding: @@ -35,6 +34,7 @@ embedding: # Semantic ranker ranker: use_ranker: true + rerank_multiplier: 8 k_ranker: 10 region: "" account_name: "" From 0aef3f262d7b0f080c4f84e3229881a13dff4245 Mon Sep 17 00:00:00 2001 From: backurs Date: Tue, 21 Apr 2026 13:19:52 -0700 Subject: [PATCH 3/6] copilot comments --- dynamic_retriever.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/dynamic_retriever.py b/dynamic_retriever.py index e0f6736..d75639a 100644 --- a/dynamic_retriever.py +++ b/dynamic_retriever.py @@ -24,7 +24,11 @@ async def vec_search(container, emb, top_k, ef): 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}])] +_SAFE_FIELD_RE = re.compile(r'^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)*$') + async def ft_field(container, field, query, top_k): + if not _SAFE_FIELD_RE.match(field): + raise ValueError(f"Invalid fulltext field name: {field!r}") terms = [t for t in re.findall(r"\w+", query) if t.lower() not in STOPWORDS and len(t) > 1] if not terms or top_k <= 0: return [] chunks = [terms[i:i+5] for i in range(0, len(terms), 5)] @@ -78,11 +82,15 @@ async def do_search(query, containers): return json.dumps([{"docid": re.search(r'^id: (.+)$', d, re.MULTILINE).group(1) if re.search(r'^id: (.+)$', d, re.MULTILINE) else "", "snippet": d[:2000]} for d in ranked]) async def do_get_doc(docid, containers): - for c in containers.values(): + 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: continue + 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): @@ -91,12 +99,16 @@ async def do_prune(docids, containers, doc_cache): if did in doc_cache: parts.append(f'\n{doc_cache[did]}\n') continue - for c in containers.values(): + 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":did}]): text = fmt(item); doc_cache[did] = text parts.append(f'\n{text}\n'); break - except: continue + 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 = [ @@ -117,7 +129,7 @@ async def process_question(q_obj, containers): 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=0, max_completion_tokens=llm_cfg["max_completion_tokens"]) + 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: {e}"); await asyncio.sleep(min(5*2**retries, 300)); continue From 114f5ae37b6e27f2d27daa6d989dfc282e2285cb Mon Sep 17 00:00:00 2001 From: backurs Date: Tue, 21 Apr 2026 17:27:57 -0700 Subject: [PATCH 4/6] copilot comments --- dynamic_retriever.py | 87 +++++++++++++++++++++++++++++++++++++------- 1 file changed, 74 insertions(+), 13 deletions(-) diff --git a/dynamic_retriever.py b/dynamic_retriever.py index d75639a..7b0c9fb 100644 --- a/dynamic_retriever.py +++ b/dynamic_retriever.py @@ -18,14 +18,20 @@ def count_tokens(msgs): # --- Search helpers --- async def embed(text): r = await embed_client.embeddings.create(input=[text], model=embed_cfg["embed_model"]) - return [float(x) for x in r.data[0].embedding[:embed_cfg.get("embed_dimensions", 1536)]] + 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}])] -_SAFE_FIELD_RE = re.compile(r'^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)*$') - async def ft_field(container, field, query, top_k): if not _SAFE_FIELD_RE.match(field): raise ValueError(f"Invalid fulltext field name: {field!r}") @@ -99,11 +105,17 @@ async def do_prune(docids, containers, doc_cache): 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'); break + text = fmt(item) + doc_cache[did] = text + parts.append(f'\n{text}\n') + found = True + break except asyncio.CancelledError: raise except Exception as e: @@ -132,7 +144,14 @@ async def process_question(q_obj, containers): 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: {e}"); await asyncio.sleep(min(5*2**retries, 300)); continue + 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: @@ -145,8 +164,34 @@ async def process_question(q_obj, containers): "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): - a = json.loads(t.function.arguments) + 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: @@ -174,15 +219,22 @@ async def _exec(t): 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(json.loads(t.function.arguments)['docids'])} docs, context reset") + 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}) - a = json.loads(t.function.arguments) - print(f" [{t.function.name}] {list(a.values())[0][:80] if isinstance(list(a.values())[0], str) else '...'}") + 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}") @@ -192,7 +244,9 @@ async def _exec(t): async def main(): questions = json.loads(Path(cfg["paths"]["questions_path"]).read_text()) - cosmos = CosmosClient(cosmos_cfg["uri"], credential=AsyncAzureCliCredential()) + use_rbac_auth = cosmos_cfg.get("use_rbac_auth", False) + credential = AsyncAzureCliCredential() if use_rbac_auth else cosmos_cfg["key"] + cosmos = CosmosClient(cosmos_cfg["uri"], credential=credential) db = cosmos.get_database_client(cosmos_cfg["database_name"]) containers = {s["id"]: db.get_container_client(s["container_name"]) for s in sources} @@ -226,7 +280,9 @@ async def main(): 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_client = AsyncAzureOpenAI(api_version=embed_cfg["api_version"], azure_endpoint=embed_cfg["embed_endpoint"], api_key=embed_cfg["embed_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"] @@ -234,7 +290,12 @@ async def main(): 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"] - _r_tok = Path(rcfg["access_token_path"]).read_text().strip() if rcfg["read_token_from_path"] else __import__("azure.identity",fromlist=["AzureCliCredential"]).AzureCliCredential(tenant_id=rcfg["tenant_id"]).get_token(rcfg["token_scope"]).token + 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) From 19370d33e0ebbd2e6bae71d3b4e809d29cfd0623 Mon Sep 17 00:00:00 2001 From: backurs Date: Tue, 21 Apr 2026 19:00:31 -0700 Subject: [PATCH 5/6] copilot comments --- dynamic_retriever.py | 39 +++++++++++++++++++++++++++------------ 1 file changed, 27 insertions(+), 12 deletions(-) diff --git a/dynamic_retriever.py b/dynamic_retriever.py index 7b0c9fb..7bb55f2 100644 --- a/dynamic_retriever.py +++ b/dynamic_retriever.py @@ -84,8 +84,12 @@ async def do_search(query, containers): 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()) - ranked = await rerank(query, [fmt(d) for d in all_d], total_k) - return json.dumps([{"docid": re.search(r'^id: (.+)$', d, re.MULTILINE).group(1) if re.search(r'^id: (.+)$', d, re.MULTILINE) else "", "snippet": d[:2000]} for d in ranked]) + 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(): @@ -244,25 +248,36 @@ async def _exec(t): 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 cosmos_cfg["key"] - cosmos = CosmosClient(cosmos_cfg["uri"], credential=credential) + 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} - results = [] - for q in tqdm(questions): - results.append(await process_question(q, containers)) + 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}") - await cosmos.close() + 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()) From 67a06f97ec3b7ae02233f4bf2bac1e7ff46aec01 Mon Sep 17 00:00:00 2001 From: backurs Date: Wed, 22 Apr 2026 12:09:41 -0700 Subject: [PATCH 6/6] refactoring and addressing comments --- dynamic_retriever.py | 47 ++++---------- utils/cosmos_retriever.py | 132 +++++--------------------------------- utils/fulltext.py | 69 ++++++++++++++++++++ utils/ranker.py | 57 ++++++++++++++++ 4 files changed, 155 insertions(+), 150 deletions(-) create mode 100644 utils/fulltext.py create mode 100644 utils/ranker.py diff --git a/dynamic_retriever.py b/dynamic_retriever.py index 7bb55f2..da7bcd5 100644 --- a/dynamic_retriever.py +++ b/dynamic_retriever.py @@ -8,8 +8,8 @@ from azure.identity import AzureCliCredential, get_bearer_token_provider from azure.identity.aio import AzureCliCredential as AsyncAzureCliCredential from azure.cosmos.aio import CosmosClient - -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"} +from utils.fulltext import fulltext_search +from utils.ranker import rerank_documents _enc = tiktoken.get_encoding("o200k_base") def count_tokens(msgs): @@ -32,37 +32,12 @@ async def vec_search(container, emb, top_k, ef): 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 ft_field(container, field, query, top_k): - if not _SAFE_FIELD_RE.match(field): - raise ValueError(f"Invalid fulltext field name: {field!r}") - terms = [t for t in re.findall(r"\w+", query) if t.lower() not in STOPWORDS and len(t) > 1] - if not terms or top_k <= 0: return [] - chunks = [terms[i:i+5] for i in range(0, len(terms), 5)] - exprs = [f'FullTextScore(c.{field}, {", ".join(chr(34)+t.replace(chr(34),"")+chr(34) for t in ch)})' for ch in chunks] - order = f"ORDER BY RANK {exprs[0]}" if len(exprs)==1 else f"ORDER BY RANK RRF({', '.join(exprs)})" - try: return [item async for item in container.query_items(query=f"SELECT TOP {top_k} * FROM c {order}", parameters=[])] - except: return [] - -async def ft_search(container, fields, query, top_k): - if not fields or top_k <= 0: return [] - if len(fields) == 1: return await ft_field(container, fields[0], query, top_k) - pf = await asyncio.gather(*(ft_field(container, f, query, top_k) for f in fields)) - sc, dm = {}, {} - for items in pf: - for r, it in enumerate(items): - d = it.get("id","") - if d: sc[d] = sc.get(d,0) + 1/(60+r+1); dm.setdefault(d, it) - return [dm[d] for d in sorted(sc, key=sc.get, reverse=True)[:top_k]] - async def rerank(query, docs, top_k): if not USE_RANKER or not docs: return docs[:top_k] - body = {"query": query, "documents": docs, "return_documents": False, "top_k": top_k, "batch_size": _r_bs} - for att in range(_r_mr): - resp = await _r_http.post(_r_url, headers=_r_hdr, json=body) - if resp.status_code in (429,502,503) and att+1 < _r_mr: await asyncio.sleep(2**att); continue - resp.raise_for_status() - return [docs[s["index"]] for s in resp.json().get("Scores",[])[:top_k] if s["index"] < len(docs)] - return docs + 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 @@ -76,7 +51,7 @@ async def do_search(query, 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(ft_search(containers[sid], fields, query, _source_cfg[sid]["fulltext_search_k"]*RERANK_MUL)) + 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: @@ -201,13 +176,17 @@ async def _exec(t): try: for h in json.loads(out): if h.get("snippet"): doc_cache[h["docid"]] = h["snippet"] - except: pass + 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: pass + 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 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