Initial commit

.gitignore

@@ -205,3 +205,8 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+
+
+data/*.json
+.vscode/
+rag_index_json/

.python-version

@@ -0,0 +1 @@
+3.14

README.md

@@ -1 +1,11 @@
-# llm-agent
+# llm-agents
+
+Install Ollama and pull `nomic-embed-text` and `llama3.1` models
+
+### Agent for NeurIPS 2025 papers
+
+Download https://neurips.cc/static/virtual/data/neurips-2025-orals-posters.json and put the file in the `./data` folder.
+
+Run `uv run build_json_rag.py` to build the vector database for RAG. 
+
+Then run using `uv run main.py`.

build_json_rag.py

@@ -0,0 +1,67 @@
+# rag_build_json.py
+from __future__ import annotations
+from dataclasses import dataclass
+import time
+from typing import List, Dict, Any, Tuple
+import json, os, math
+from pathlib import Path
+
+import numpy as np
+import faiss
+from litellm import embedding
+from tqdm import tqdm
+
+API_BASE = "http://localhost:11434"
+EMBED_MODEL = "ollama/nomic-embed-text"   # pull with: ollama pull nomic-embed-text
+INDEX_DIR = "rag_index_json"
+
+
+def _load_papers(json_path: str) -> List[Paper]:
+    data = json.loads(Path(json_path).read_text())
+    papers = data["results"]
+    return papers
+
+def _embed(texts: List[str], batch: int = 1) -> np.ndarray:
+    vecs: List[List[float]] = []
+    for i in tqdm(range(0, len(texts), batch)):
+        chunk = texts[i:i+batch]
+        retry = 0
+        max_retries = 100
+        while retry < max_retries:
+            try:
+                resp = embedding(model=EMBED_MODEL, input=chunk, api_base=API_BASE)
+                break
+            except Exception as e:
+                print(f"Error during embedding batch {i}-{i+batch}: {e}")
+                # print(chunk)
+                retry += 1
+                if retry >= max_retries:
+                    raise Exception("Max retries reached for embedding.")
+                print(f"Retrying {retry} ...")
+                time.sleep(0.1)
+
+        vecs.extend([d["embedding"] for d in resp["data"]])
+        # time.sleep(5)  # avoid rate limit
+    arr = np.array(vecs, dtype="float32")
+    # cosine similarity: normalize to unit length and use IndexFlatIP
+    norms = np.linalg.norm(arr, axis=1, keepdims=True) + 1e-12
+    return arr / norms
+
+def build_index(json_path: str, out_dir: str = INDEX_DIR) -> None:
+    out = Path(out_dir); out.mkdir(parents=True, exist_ok=True)
+    papers = _load_papers(json_path)
+
+    texts = [f"{p['name']}\n\n{p['abstract']}" if "name" in p and "abstract" in p else "" for p in papers]
+    X = _embed(texts)
+    dim = X.shape[1]
+    index = faiss.IndexFlatIP(dim)
+    index.add(X)
+
+    faiss.write_index(index, str(out / "faiss.index"))
+    # metadata alongside embeddings
+    (out / "meta.json").write_text(json.dumps(papers, ensure_ascii=False, indent=4, sort_keys=True))
+    print(f"Indexed {len(papers)} papers → {out.resolve()}")
+
+if __name__ == "__main__":
+    # build_index("data/papers_test.json")
+    build_index("data/neurips-2025-orals-posters-pretty.json")

main.py

@@ -0,0 +1,16 @@
+from tools.rag_search import search_papers  # <- the tool we expose
+from utils.tool_chat import ToolChat
+import sys
+
+if __name__ == "__main__":
+    chat = ToolChat()
+    system = {"role": "system", "content": "You are a research assistant. Use the search tool when helpful, but you may want to search for a larger number of relevant papers, then select from those that are most relevant to the user's query by further examining the titles and abstracts of the papers found. Cite titles, abstracts, and OpenReview URLs in your answers."}
+    user = {"role": "user", "content": "Find papers on large language models. Give me the title, abstract, and OpenReview URLs."}
+
+    if len(sys.argv) > 1:
+        user["content"] = " ".join(sys.argv[1:])
+    else:
+        user["content"] = input("Enter user prompt: ")
+
+    resp = chat.tool_loop([system, user], registry={"search_papers": search_papers})
+    print(resp["choices"][0]["message"]["content"])

pyproject.toml

@@ -0,0 +1,16 @@
+[project]
+name = "llm-agents"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.14"
+dependencies = [
+    "faiss-cpu>=1.12.0",
+    "httpx>=0.28.1",
+    "litellm>=1.79.0",
+    "pydantic>=2.12.3",
+    "pydantic-settings>=2.11.0",
+    "rich>=14.2.0",
+    "sentence-transformers>=5.1.2",
+    "typer>=0.20.0",
+]

tools/rag_search.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+from typing import List, Dict, Any
+import json
+from pathlib import Path
+import numpy as np
+import faiss
+from litellm import embedding
+import time
+
+API_BASE = "http://localhost:11434"
+EMBED_MODEL = "ollama/nomic-embed-text"
+INDEX_DIR = "rag_index_json"
+
+# Load once at import (fast)
+_FAISS = faiss.read_index(str(Path(INDEX_DIR) / "faiss.index"))
+_META = json.loads(Path(INDEX_DIR, "meta.json").read_text())
+
+def _embed_query(q: str) -> np.ndarray:
+    resp = embedding(model=EMBED_MODEL, input=[q], api_base=API_BASE)
+    v = np.array(resp["data"][0]["embedding"], dtype="float32")
+    v = v / (np.linalg.norm(v) + 1e-12)
+    return v.reshape(1, -1)
+
+def _snippet(text: str, max_chars: int = 240) -> str:
+    text = " ".join(text.split())
+    return text if len(text) <= max_chars else text[:max_chars-1] + "…"
+
+def search_papers(query: str, k: int = 10) -> Dict[str, Any]:
+    """
+    Search top-k relevant papers over title+abstract. You need to specify k as the total number of results that you would like to get.
+    Returns JSON with 'query' and 'results' (the results include metadata of the paper).
+    """
+    try:
+        k = int(k)
+    except Exception:
+        k = 5
+    
+    start_time = time.time()
+    qv = _embed_query(query)
+    scores, idxs = _FAISS.search(qv, k)
+    scores, idxs = scores[0].tolist(), idxs[0].tolist()
+    print(f"Search for '{query}' took {time.time() - start_time:.3f}s")
+
+    results = []
+    for score, i in zip(scores, idxs):
+        if i == -1:
+            continue
+        m = _META[i]
+        text = f"{m['name']}\n\n{m['abstract']}"
+        results.append(m)
+    return {"query": query, "results": results}

utils/tool_chat.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+from typing import Any, Dict, List, Callable, get_args, get_origin, Literal, Annotated, Union
+from dataclasses import dataclass, field
+import json, inspect
+from litellm import completion
+
+from tools.rag_search import search_papers  # <- the tool we expose
+
+# ---- Minimal auto-schema from function signature ----
+def _json_type(t: Any) -> Dict[str, Any]:
+    origin, args = get_origin(t), get_args(t)
+    if origin is Literal:
+        return {"type": "string", "enum": list(args)}
+    if origin in (list, List):
+        return {"type": "array", "items": {"type": "string"}}
+    if t in (str,): return {"type": "string"}
+    if t in (int,): return {"type": "integer"}
+    if t in (float,): return {"type": "number"}
+    if t in (bool,): return {"type": "boolean"}
+    return {"type": "string"}
+
+def infer_tool(func: Callable[..., Any]) -> Dict[str, Any]:
+    sig = inspect.signature(func)
+    hints = getattr(func, "__annotations__", {})
+    props, required = {}, []
+    for name, p in sig.parameters.items():
+        if name in ("self", "cls"): continue
+        schema = _json_type(hints.get(name, str))
+        if p.default is inspect._empty: required.append(name)
+        props[name] = schema
+    return {
+        "type": "function",
+        "function": {
+            "name": func.__name__,
+            "description": (inspect.getdoc(func) or f"Call {func.__name__}"),
+            "parameters": {"type": "object", "properties": props, "required": required},
+        },
+    }
+
+# ---- LiteLLM chat wrapper with tool loop ----
+@dataclass
+class ToolChat:
+    model: str = "ollama_chat/llama3.1"
+    # model: str = "ollama_chat/ibm/granite4:350m"
+    api_base: str = "http://localhost:11434"
+    default_params: Dict[str, Any] = field(default_factory=lambda: {"temperature": 0.2, "max_tokens": 2000})
+
+    def tool_loop(self, messages: List[Dict[str, Any]], registry: Dict[str, Callable[..., Any]], max_rounds: int = 3) -> Dict[str, Any]:
+        tools = [infer_tool(fn) for fn in registry.values()]
+        msgs = list(messages)
+        for _ in range(max_rounds):
+            resp = completion(model=self.model, messages=msgs, tools=tools, tool_choice="auto", api_base=self.api_base, **self.default_params)
+
+            msg = resp["choices"][0].get("message", {})
+            calls = msg.get("tool_calls") or []
+            print("tool_calls:", calls)
+
+            if not calls:
+                return resp
+
+            # execute tools and append results
+            for call in calls:
+                name = call["function"]["name"]
+                args = call["function"].get("arguments", "{}")
+                try:
+                    parsed = json.loads(args) if isinstance(args, str) else (args or {})
+                except json.JSONDecodeError:
+                    parsed = {}
+                out = registry[name](**parsed)
+                msgs.append({
+                    "role": "tool",
+                    "tool_call_id": call.get("id"),
+                    "name": name,
+                    "content": json.dumps(out, ensure_ascii=False),
+                })
+        return completion(model=self.model, messages=msgs, api_base=self.api_base, **self.default_params)