Extending Cognotik

// Example: Creating a custom DynamicEnum
class MyExtension(name: String) : DynamicEnum<MyExtension>(name) {
    companion object {
        val Default = MyExtension("Default")

        init {
            register(MyExtension::class.java, Default)
        }

        fun values() = values(MyExtension::class.java)
        fun valueOf(name: String) = valueOf(MyExtension::class.java, name)
    }
}

val CustomProvider = object : APIProvider("CustomProvider", "https://api.custom.com") {
    override fun getChatModels(key: SecureString, baseUrl: String): List<ChatModel> {
        // Return available chat models
        return listOf(/* your models */)
    }

    override fun getChatClient(
        key: SecureString,
        base: String,
        workPool: ExecutorService,
        logLevel: Level,
        logStreams: MutableList<BufferedOutputStream>,
        scheduledPool: ListeningScheduledExecutorService
    ): ChatClientInterface {
        return YourChatClient(key, base, workPool, scheduledPool)
    }

    // Optional: Override for custom authorization
    override fun authorize(request: HttpRequest, key: String, apiBase: String) {
        request.addHeader("X-Custom-Auth", key)
    }

    // Optional: Support embeddings
    override fun getEmbeddingModels(key: SecureString, baseUrl: String): List<EmbeddingModel> {
        return listOf(/* your embedding models */)
    }

    // Optional: Support image generation
    override fun getImageModels(key: SecureString, baseUrl: String): List<ImageModel> {
        return listOf(/* your image models */)
    }
}

// Register the provider
APIProvider.register(APIProvider::class.java, CustomProvider)

val CustomTool = object : ToolProvider("CustomTool") {
    override fun getExecutables() = listOf("custom-tool", "custom-tool.exe")

    override fun getVersion(path: String): String? {
        return runCommand(listOf(path, "--version"))
    }

    // Optional: Custom validation
    override fun validate(path: String): Boolean {
        return try {
            val version = getVersion(path)
            version != null && version.contains("CustomTool")
        } catch (e: Exception) {
            false
        }
    }
}

// Register the tool
ToolProvider.register(ToolProvider::class.java, CustomTool)

// Discover all tools from PATH
val tools = ToolProvider.discoverAllToolsFromPath()

// Scan a directory recursively
val projectTools = ToolProvider.scanRecursive(File("/path/to/project"), depth = 3)

// Resolve a specific tool
val pythonPath = ToolProvider.Python.resolve("/usr/local")

val CustomRuntime = CodeRuntimes("CustomRuntime", "Execute custom scripts", "custom")

CodeRuntimes.registerConstructor(CustomRuntime) { defs ->
    ProcessCodeRuntime(
        timeoutMinutes = defs["timeoutMinutes"]?.toString()?.toLongOrNull() ?: 15L,
        workingDir = defs["workingDir"]?.toString()?.let { File(it) } ?: File("."),
        env = defs["env"]?.let { it as Map<String, String> },
        lang = "custom",
        command = listOf("custom-interpreter")
    )
}

val runtime = CodeRuntimes.getRuntime(
    CodeRuntimes.PythonRuntime,
    mapOf(
        "timeoutMinutes" to 30L,
        "workingDir" to "/path/to/project"
    )
)

val result = runtime.execute("print('Hello, World!')")

// 1. Define the config class
class CustomModeConfig : CognitiveModeConfig() {
    var customSetting: String = "default"
}

// 2. Define the mode type
val CustomMode = CognitiveModeType(
    "CustomMode",
    CustomModeConfig::class.java,
    inputCnt = 1
)

// 3. Implement the mode
class CustomCognitiveMode(
    config: OrchestrationConfig,
    session: Session,
    user: User
) : CognitiveMode<CustomModeConfig>(config, session, user) {
    override fun run(messages: List<String>): String {
        // Implement your cognitive mode logic
        return "Result"
    }
}

// 4. Register (done in CognitiveModeType companion object)

// 1. Define your state class
data class CustomState(
    val objective: String? = null,
    val progress: MutableList<String>? = null,
    val metrics: MutableMap<String, Double>? = null
)

// 2. Implement the strategy
class CustomStrategy : CognitiveSchemaStrategy(
    "Custom Strategy",
    "Description of your strategy"
) {
    override fun initialize(
        userMessage: String,
        contextData: List<String>,
        orchestrationConfig: OrchestrationConfig,
        task: SessionTask,
        describer: TaskContextYamlDescriber
    ): Any {
        return ParsedAgent(
            name = "CustomInitializer",
            resultClass = CustomState::class.java,
            exampleInstance = CustomState(
                objective = "Example objective",
                progress = mutableListOf("Step 1"),
                metrics = mutableMapOf("completion" to 0.0)
            ),
            prompt = "Initialize the custom state based on the user request...",
            model = orchestrationConfig.defaultSmart.getChildClient(task),
            parsingChatter = orchestrationConfig.defaultFast.getChildClient(task),
            temperature = orchestrationConfig.temperature,
            describer = describer
        ).answer(listOf(userMessage) + contextData).obj
    }

    override fun update(
        currentState: Any,
        completedTasks: List<AdaptivePlanningMode.ExecutionRecord>,
        userMessage: String?,
        contextData: List<String>,
        orchestrationConfig: OrchestrationConfig,
        task: SessionTask,
        describer: TaskContextYamlDescriber
    ): Any {
        // Update logic
        return currentState
    }

    override fun formatState(state: Any): String {
        return JsonUtil.toJson(state)
    }

    override fun getTaskSelectionGuidance(state: Any): String {
        return "Select tasks based on current metrics and progress."
    }
}

// 3. Register the strategy
DynamicEnum.register(CognitiveSchemaStrategy::class.java, CustomStrategy())

// 1. Define the execution config
data class CustomTaskConfig(
    override var task_type: String? = CustomTask.name,
    override var task_description: String? = null,
    var customParameter: String? = null
) : TaskExecutionConfig()

// 2. Define the task settings
data class CustomTaskSettings(
    override var task_type: String? = CustomTask.name,
    override var enabled: Boolean = true,
    var defaultValue: String = "default"
) : TaskTypeConfig()

// 3. Implement the task
class CustomTaskImpl(
    settings: OrchestrationConfig,
    task: CustomTaskConfig?
) : AbstractTask<CustomTaskConfig, CustomTaskSettings>(settings, task) {

    override fun promptSegment(): String {
        return """
            CustomTask - Performs custom operations
            ** Specify the custom parameter
        """.trimIndent()
    }

    override fun run(
        agent: PlanCoordinator,
        messages: List<String>,
        task: SessionTask,
        api: API,
        resultFn: (String) -> Unit,
        api2: OpenAIClient,
        planSettings: PlanSettings
    ): TaskResult {
        // Implementation
        return TaskResult("Custom task completed")
    }

    companion object {
        val CustomTask = TaskType(
            name = "CustomTask",
            category = "Custom",
            taskClass = CustomTaskImpl::class.java,
            executionConfigClass = CustomTaskConfig::class.java,
            taskSettingsClass = CustomTaskSettings::class.java,
            description = "Performs custom operations"
        )
    }
}

// 4. Register (done via registerConstructor in TaskType companion object)

// Custom serializer (usually auto-generated)
class MyExtensionSerializer : DynamicEnumSerializer<MyExtension>(MyExtension::class.java)
class MyExtensionDeserializer : DynamicEnumDeserializer<MyExtension>(MyExtension::class.java)

// Apply to your class
@JsonDeserialize(using = MyExtensionDeserializer::class)
@JsonSerialize(using = MyExtensionSerializer::class)
class MyExtension(name: String) : DynamicEnum<MyExtension>(name)