refactored UI function config access to use a hook / helper (tensorzero#2910)

* refactored UI function config access to use a hook / helper

* made hook / helper for getting all functions

---------

Co-authored-by: Gabriel Bianconi <1275491+GabrielBianconi@users.noreply.github.com>

Author: virajmehta

ui/app/components/function/FunctionInfo.tsx

@@ -1,5 +1,5 @@
 import { Code } from "~/components/ui/code";
-import { useConfig } from "~/context/config";
+import { useFunctionConfig } from "~/context/config";
 import { AlertDialog } from "~/components/ui/AlertDialog";
 import type { FunctionConfig } from "tensorzero-node";
 import { FunctionLink } from "./FunctionLink";
@@ -13,8 +13,7 @@ export function FunctionInfo({
   functionName,
   functionType,
 }: FunctionInfoProps) {
-  const config = useConfig();
-  const functionConfig = config.functions[functionName];
+  const functionConfig = useFunctionConfig(functionName);
 
   return (
     <>

ui/app/components/function/FunctionLink.tsx

@@ -1,6 +1,6 @@
 import { Link } from "react-router";
 import { AlertDialog } from "~/components/ui/AlertDialog";
-import { useConfig } from "~/context/config";
+import { useFunctionConfig } from "~/context/config";
 import type { ReactNode } from "react";
 
 type FunctionLinkProps = {
@@ -9,8 +9,7 @@ type FunctionLinkProps = {
 };
 
 export function FunctionLink({ functionName, children }: FunctionLinkProps) {
-  const config = useConfig();
-  const functionConfig = config.functions[functionName];
+  const functionConfig = useFunctionConfig(functionName);
   return functionConfig ? (
     <Link to={`/observability/functions/${functionName}`}>{children}</Link>
   ) : (

ui/app/components/function/variant/VariantLink.tsx

@@ -1,6 +1,6 @@
 import { Link } from "react-router";
 import { AlertDialog } from "~/components/ui/AlertDialog";
-import { useConfig } from "~/context/config";
+import { useFunctionConfig } from "~/context/config";
 import type { ReactNode } from "react";
 
 type VariantLinkProps = {
@@ -14,8 +14,7 @@ export function VariantLink({
   functionName,
   children,
 }: VariantLinkProps) {
-  const config = useConfig();
-  const functionConfig = config.functions[functionName];
+  const functionConfig = useFunctionConfig(functionName);
   const variantConfig = functionConfig?.variants[variantName];
   return variantConfig ? (
     <Link

ui/app/components/ui/TableItems.tsx

@@ -7,7 +7,7 @@ import {
 } from "./tooltip";
 import { getFunctionTypeIcon } from "~/utils/icon";
 import { formatDate } from "~/utils/date";
-import { useConfig } from "~/context/config";
+import { useFunctionConfig } from "~/context/config";
 import { AlertDialog } from "~/components/ui/AlertDialog";
 
 interface TableItemShortUuidProps {
@@ -77,9 +77,8 @@ function TableItemFunction({
   functionType,
   link,
 }: TableItemFunctionProps) {
-  const config = useConfig();
   const functionIconConfig = getFunctionTypeIcon(functionType);
-  const functionConfig = config.functions[functionName];
+  const functionConfig = useFunctionConfig(functionName);
 
   const baseClasses =
     "flex items-center text-sm text-fg-primary gap-2 rounded-md font-mono";

ui/app/context/config.tsx

@@ -21,6 +21,29 @@ export function useConfig() {
   return config;
 }
 
+/**
+ * Hook to get a specific function configuration by name
+ * @param functionName - The name of the function to retrieve
+ * @returns The function configuration object or null if not found
+ */
+export function useFunctionConfig(functionName: string | null) {
+  const config = useConfig();
+  if (!functionName) {
+    return null;
+  }
+  // eslint-disable-next-line no-restricted-syntax
+  return config.functions[functionName] || null;
+}
+/**
+ * Hook to get all function configs
+ * @returns The function configuration object or null if not found
+ */
+export function useAllFunctionConfigs() {
+  const config = useConfig();
+  // eslint-disable-next-line no-restricted-syntax
+  return config.functions;
+}
+
 export function ConfigProvider({
   children,
   value,

ui/app/routes/api/curated_inferences/count.route.ts

@@ -6,7 +6,7 @@ import {
 } from "~/utils/clickhouse/curation.server";
 import { countInferencesForFunction } from "~/utils/clickhouse/inference.server";
 import { getFeedbackConfig } from "~/utils/config/feedback";
-import { getConfig } from "~/utils/config/index.server";
+import { getConfig, getFunctionConfig } from "~/utils/config/index.server";
 
 /// Count the number of inferences, feedbacks, and curated inferences for a given function and metric
 /// This is used to determine the number of inferences to display in the UI
@@ -23,7 +23,9 @@ export async function loader({
   const threshold = parseFloat(url.searchParams.get("threshold") || "0");
 
   const config = await getConfig();
-  const functionConfig = config.functions[functionName || ""];
+  const functionConfig = functionName
+    ? await getFunctionConfig(functionName, config)
+    : null;
   if (functionName && !functionConfig) {
     throw data(`Function ${functionName} not found in config`, { status: 404 });
   }

ui/app/routes/api/function/$function_name/feedback_counts.route.ts

@@ -1,5 +1,5 @@
 import { data, type LoaderFunctionArgs } from "react-router";
-import { getConfig } from "~/utils/config/index.server";
+import { getConfig, getFunctionConfig } from "~/utils/config/index.server";
 import { getInferenceTableName } from "~/utils/clickhouse/common";
 import {
   queryMetricsWithFeedback,
@@ -18,7 +18,7 @@ export async function loader({
 
   try {
     const config = await getConfig();
-    const functionConfig = config.functions[functionName];
+    const functionConfig = await getFunctionConfig(functionName, config);
     if (!functionConfig) {
       throw data(`Function ${functionName} not found in config`, {
         status: 404,

ui/app/routes/datasets/$dataset_name/datapoint/$id/DatapointBasicInfo.tsx

@@ -1,4 +1,4 @@
-import { useConfig } from "~/context/config";
+import { useFunctionConfig } from "~/context/config";
 import type { ParsedDatasetRow } from "~/utils/clickhouse/datasets";
 import {
   BasicInfoLayout,
@@ -30,8 +30,7 @@ const createTimestampTooltip = (timestamp: string | number | Date) => {
 };
 
 export default function DatapointBasicInfo({ datapoint }: BasicInfoProps) {
-  const config = useConfig();
-  const function_config = config.functions[datapoint.function_name];
+  const function_config = useFunctionConfig(datapoint.function_name);
   const type = function_config?.type;
   if (!type) {
     throw new Error(`Function ${datapoint.function_name} not found`);

ui/app/routes/datasets/$dataset_name/datapoint/$id/route.tsx

@@ -21,7 +21,7 @@ import {
   SectionsGroup,
 } from "~/components/layout/PageLayout";
 import { Badge } from "~/components/ui/badge";
-import { useConfig } from "~/context/config";
+import { useFunctionConfig } from "~/context/config";
 import { resolvedInputToTensorZeroInput } from "~/routes/api/tensorzero/inference";
 import {
   prepareInferenceActionRequest,
@@ -37,7 +37,7 @@ import {
   getDatasetCounts,
   staleDatapoint,
 } from "~/utils/clickhouse/datasets.server";
-import { getConfig } from "~/utils/config/index.server";
+import { getConfig, getFunctionConfig } from "~/utils/config/index.server";
 import { logger } from "~/utils/logger";
 import { getTensorZeroClient } from "~/utils/tensorzero.server";
 import type { Route } from "./+types/route";
@@ -77,7 +77,10 @@ export async function action({ request }: ActionFunctionArgs) {
   const parsedFormData: ParsedDatasetRow =
     ParsedDatasetRowSchema.parse(cleanedData);
   const config = await getConfig();
-  const functionConfig = config.functions[parsedFormData.function_name];
+  const functionConfig = await getFunctionConfig(
+    parsedFormData.function_name,
+    config,
+  );
   if (!functionConfig) {
     return new Response(
       `Failed to find function config for function ${parsedFormData.function_name}`,
@@ -198,7 +201,6 @@ export default function DatapointPage({ loaderData }: Route.ComponentProps) {
   const [output, setOutput] = useState<typeof datapoint.output>(
     datapoint.output,
   );
-  const config = useConfig();
   const [isEditing, setIsEditing] = useState(false);
 
   const canSave = useMemo(() => {
@@ -259,9 +261,8 @@ export default function DatapointPage({ loaderData }: Route.ComponentProps) {
     }
   };
 
-  const variants = Object.keys(
-    config.functions[datapoint.function_name]?.variants || {},
-  );
+  const functionConfig = useFunctionConfig(datapoint.function_name);
+  const variants = Object.keys(functionConfig?.variants || {});
 
   const variantInferenceFetcher = useInferenceActionFetcher();
   const variantSource = "datapoint";

ui/app/routes/datasets/builder/DatasetBuilderForm.tsx

@@ -6,7 +6,11 @@ import {
 } from "./types";
 import { FunctionFormField } from "~/components/function/FunctionFormField";
 import { DatasetFormField } from "~/components/dataset/DatasetFormField";
-import { useConfig } from "~/context/config";
+import {
+  useAllFunctionConfigs,
+  useConfig,
+  useFunctionConfig,
+} from "~/context/config";
 import CurationMetricSelector from "~/components/metric/CurationMetricSelector";
 import { useCountFetcher } from "~/routes/api/curated_inferences/count.route";
 import { useFetcher } from "react-router";
@@ -54,14 +58,16 @@ export function DatasetBuilderForm() {
     metricName: metricName ?? undefined,
     threshold: threshold ?? undefined,
   });
+  const functionConfig = useFunctionConfig(functionName ?? "");
+
   useEffect(() => {
     const metricConfig = config.metrics[metricName ?? ""];
     form.setValue("metric_config", metricConfig ? metricConfig : undefined);
-    const functionType = config.functions[functionName ?? ""]?.type;
+    const functionType = functionConfig?.type;
     if (functionType) {
       form.setValue("type", functionType);
     }
-  }, [metricName, functionName, config, form]);
+  }, [metricName, functionName, config, form, functionConfig]);
 
   // Handle form submission response
   useEffect(() => {
@@ -132,7 +138,7 @@ export function DatasetBuilderForm() {
           <FunctionFormField
             control={form.control}
             name="function"
-            functions={config.functions}
+            functions={useAllFunctionConfigs()}
             onSelect={() => {
               form.resetField("variant");
             }}