@@ -5,6 +5,7 @@ import { useListSupportedModelsQuery } from "../query";
55import { useConversationUiStore } from "../stores/conversation/conversation-ui-store" ;
66
77export type Model = {
8+ id ?: string ;
89 name : string ;
910 slug : string ;
1011 provider : string ;
@@ -39,6 +40,7 @@ const fallbackModels: Model[] = [
3940] ;
4041
4142const mapSupportedModelToModel = ( supportedModel : SupportedModel ) : Model => ( {
43+ id : supportedModel . id || undefined ,
4244 name : supportedModel . name ,
4345 slug : supportedModel . slug ,
4446 provider : extractProvider ( supportedModel . slug ) ,
@@ -53,7 +55,7 @@ const mapSupportedModelToModel = (supportedModel: SupportedModel): Model => ({
5355
5456export const useLanguageModels = ( ) => {
5557 const { currentConversation, setCurrentConversation } = useConversationStore ( ) ;
56- const { setLastUsedModelSlug } = useConversationUiStore ( ) ;
58+ const { lastUsedCustomModelId , setLastUsedModelSlug, setLastUsedCustomModelId } = useConversationUiStore ( ) ;
5759 const { data : supportedModelsResponse } = useListSupportedModelsQuery ( ) ;
5860
5961 const models : Model [ ] = useMemo ( ( ) => {
@@ -64,19 +66,25 @@ export const useLanguageModels = () => {
6466 } , [ supportedModelsResponse ] ) ;
6567
6668 const currentModel = useMemo ( ( ) => {
67- const model = models . find ( ( m ) => m . slug === currentConversation . modelSlug ) ;
69+ if ( lastUsedCustomModelId ) {
70+ const customModel = models . find ( ( m ) => m . isCustom && m . id === lastUsedCustomModelId ) ;
71+ if ( customModel ) return customModel ;
72+ }
73+
74+ const model = models . find ( ( m ) => ! m . isCustom && m . slug === currentConversation . modelSlug ) ;
6875 return model || models [ 0 ] ;
69- } , [ models , currentConversation . modelSlug ] ) ;
76+ } , [ models , currentConversation . modelSlug , lastUsedCustomModelId ] ) ;
7077
7178 const setModel = useCallback (
7279 ( model : Model ) => {
7380 setLastUsedModelSlug ( model . slug ) ;
81+ setLastUsedCustomModelId ( model . isCustom ? ( model . id ?? "" ) : "" ) ;
7482 setCurrentConversation ( {
7583 ...currentConversation ,
7684 modelSlug : model . slug ,
7785 } ) ;
7886 } ,
79- [ setCurrentConversation , currentConversation , setLastUsedModelSlug ] ,
87+ [ setCurrentConversation , currentConversation , setLastUsedModelSlug , setLastUsedCustomModelId ] ,
8088 ) ;
8189
8290 return { models, currentModel, setModel } ;
0 commit comments