hailey.at
An MCP server for Osprey
1from pydantic import BaseModel
2
3
4class UdfArgumentSpec(BaseModel):
5 name: str
6 type: str
7 default: str | None = None
8 doc: str | None = None
9
10
11class Udf(BaseModel):
12 name: str
13 return_type: str
14 argument_specs: list[UdfArgumentSpec] = []
15 doc: str | None = None
16 category: str | None = None
17
18 def signature(self) -> str:
19 args = ", ".join(
20 f"{arg.name}: {arg.type}" + (f" = {arg.default}" if arg.default else "")
21 for arg in self.argument_specs
22 )
23 return f"{self.name}({args}) -> {self.return_type}"
24
25
26class UdfCategory(BaseModel):
27 name: str | None = None
28 udfs: list[Udf] = []
29
30
31class UdfCatalog(BaseModel):
32 udf_categories: list[UdfCategory] = []
33
34 def all_udfs(self) -> list[Udf]:
35 return [udf for cat in self.udf_categories for udf in cat.udfs]
36
37 def find_udf(self, name: str) -> Udf | None:
38 name_lower = name.lower()
39 for udf in self.all_udfs():
40 if udf.name.lower() == name_lower:
41 return udf
42 return None
43
44 def udfs_by_category(self, category: str) -> list[Udf]:
45 for cat in self.udf_categories:
46 if cat.name and cat.name.lower() == category.lower():
47 return cat.udfs
48 return []
49
50 def format_for_llm(self) -> str:
51 lines = ["# Available UDFs\n"]
52
53 for cat in self.udf_categories:
54 cat_name = cat.name or "Other"
55 lines.append(f"## {cat_name}\n")
56
57 for udf in cat.udfs:
58 lines.append(f"### {udf.name}")
59 lines.append("```")
60 lines.append(udf.signature())
61 lines.append("```")
62 if udf.doc:
63 lines.append(udf.doc)
64
65 if udf.argument_specs:
66 lines.append("\n**Parameters:**")
67 for arg in udf.argument_specs:
68 arg_doc = arg.doc or arg.type
69 default = f" (default: {arg.default})" if arg.default else ""
70 lines.append(f"- `{arg.name}`: {arg_doc}{default}")
71
72 lines.append("")
73
74 return "\n".join(lines)