Create embeddings.py

mcp/embeddings.py

@@ -0,0 +1,46 @@
+# ── mcp/embeddings.py ───────────────────────────────────────────────────
+import os, asyncio
+from huggingface_hub import InferenceClient
+from sklearn.cluster import KMeans
+
+# Use your HF token for pro endpoints
+HF_TOKEN = os.getenv("HF_TOKEN")
+EMBED_MODEL = "sentence-transformers/all-mpnet-base-v2"
+client = InferenceClient(token=HF_TOKEN)
+
+async def embed_texts(texts: list[str]) -> list[list[float]]:
+    """
+    Compute embeddings for a list of texts via HF Inference API.
+    """
+    def _embed(t):
+        return client.embed(model=EMBED_MODEL, inputs=t)
+    # run in threadpool
+    tasks = [asyncio.to_thread(_embed, t) for t in texts]
+    return await asyncio.gather(*tasks)
+
+async def cluster_embeddings(embs: list[list[float]], n_clusters: int = 5) -> list[int]:
+    """
+    Cluster embeddings into n_clusters, return list of cluster labels.
+    """
+    kmeans = KMeans(n_clusters=n_clusters, random_state=0)
+    return kmeans.fit_predict(embs).tolist()
+
+
+# ── mcp/protocols.py ───────────────────────────────────────────────────
+import asyncio
+from mcp.openai_utils import ai_qa
+from mcp.gemini import gemini_qa
+
+async def draft_protocol(question: str, context: str, llm: str = "openai") -> str:
+    """
+    Draft a detailed experimental protocol for a given hypothesis/question.
+    """
+    if llm.lower() == "gemini":
+        qa_fn = gemini_qa
+    else:
+        qa_fn = ai_qa
+    prompt = (
+        "You are a senior researcher. Draft a step-by-step experimental protocol to test: "
+        f"{question}\nContext:\n{context}\nInclude materials, methods, controls, expected outcomes."
+    )
+    return await qa_fn(prompt)